diff --git a/src/environments.rs b/src/environments.rs index ac8a7fa..87d4f80 100644 --- a/src/environments.rs +++ b/src/environments.rs @@ -5,6 +5,7 @@ use chrono::{DateTime, Utc}; use parking_lot::RwLock; use crate::config::settings::EnvironmentKeyPair; +use crate::identities::EnvironmentKeyType; /// A server-side (`ser.`) key together with the validity metadata the /// proxy config endpoint reports. Statically configured keys carry no @@ -32,6 +33,14 @@ pub struct EnvironmentKeys { } impl EnvironmentKeys { + pub fn key_type(&self, environment_key: &str) -> EnvironmentKeyType { + if environment_key == self.client_key { + EnvironmentKeyType::Client + } else { + EnvironmentKeyType::Server + } + } + /// 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()) diff --git a/src/identities/mod.rs b/src/identities/mod.rs new file mode 100644 index 0000000..b2f3715 --- /dev/null +++ b/src/identities/mod.rs @@ -0,0 +1,3 @@ +pub mod store; + +pub use store::{EnvironmentKeyType, IdentityStore, IdentityTraits}; diff --git a/src/identities/store.rs b/src/identities/store.rs new file mode 100644 index 0000000..6f68997 --- /dev/null +++ b/src/identities/store.rs @@ -0,0 +1,30 @@ +use crate::models::{IdentityWithTraits, TraitModel}; +use async_trait::async_trait; + +#[async_trait] +pub trait IdentityStore: Send + Sync { + async fn load_traits( + &self, + environment: &serde_json::Value, + identity: &IdentityWithTraits, + ) -> IdentityTraits; + + fn save_traits( + &self, + environment: &serde_json::Value, + identity: &IdentityWithTraits, + key_type: EnvironmentKeyType, + ); +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct IdentityTraits { + pub traits: Vec, + pub system_traits: Vec, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum EnvironmentKeyType { + Client, + Server, +} diff --git a/src/lib.rs b/src/lib.rs index c13d760..4054243 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,6 +2,7 @@ pub mod cache; pub mod config; pub mod environments; pub mod error; +pub mod identities; pub mod models; pub mod routes; pub mod server; diff --git a/src/main.rs b/src/main.rs index aa74557..f8c213b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,5 +2,5 @@ use edge_proxy::config::settings::get_settings; #[tokio::main] async fn main() -> anyhow::Result<()> { - edge_proxy::run(get_settings()?).await + edge_proxy::run(get_settings()?, None).await } diff --git a/src/routes/mod.rs b/src/routes/mod.rs index fc419b0..b00ec37 100644 --- a/src/routes/mod.rs +++ b/src/routes/mod.rs @@ -6,6 +6,7 @@ pub mod health; pub mod identities; use crate::config::AppSettings; +use crate::identities::IdentityStore; use crate::services::EnvironmentService; use axum::{ Router, @@ -15,9 +16,16 @@ use axum::{ use std::sync::Arc; use tower_http::{compression::CompressionLayer, normalize_path::NormalizePath, trace::TraceLayer}; -pub fn create_router(settings: AppSettings) -> (Router, Arc) { +pub fn create_router( + settings: AppSettings, + identity_store: Option>, +) -> (Router, Arc) { let cors = cors::layer(&settings.allow_origins); - let environment_service = Arc::new(EnvironmentService::new(settings)); + let environment_service = EnvironmentService::new(settings); + let environment_service = Arc::new(match identity_store { + Some(identity_store) => environment_service.with_identity_store(identity_store), + None => environment_service, + }); let router = Router::new() // Health check routes diff --git a/src/server.rs b/src/server.rs index 0ccba09..9c5d95d 100644 --- a/src/server.rs +++ b/src/server.rs @@ -1,12 +1,17 @@ use crate::config::AppSettings; use crate::config::logging::setup_logging; +use crate::identities::IdentityStore; use crate::routes::create_router; use std::net::{IpAddr, Ipv4Addr, SocketAddr}; +use std::sync::Arc; use tracing::{info, warn}; const DEFAULT_HOST: IpAddr = IpAddr::V4(Ipv4Addr::UNSPECIFIED); -pub async fn run(settings: AppSettings) -> anyhow::Result<()> { +pub async fn run( + settings: AppSettings, + identity_store: Option>, +) -> anyhow::Result<()> { setup_logging(&settings.logging); info!("Starting Edge Proxy server..."); @@ -27,7 +32,7 @@ pub async fn run(settings: AppSettings) -> anyhow::Result<()> { ); } - let (app, environment_service) = create_router(settings.clone()); + let (app, environment_service) = create_router(settings.clone(), identity_store); // Refreshes must never overlap: a delayed older poll finishing after a // newer one could restore removed environments or rotated keys. The diff --git a/src/services/environment.rs b/src/services/environment.rs index 5ef66dd..d6c2628 100644 --- a/src/services/environment.rs +++ b/src/services/environment.rs @@ -2,8 +2,9 @@ use crate::cache::{EnvironmentsCache, LocalMemEnvironmentsCache}; use crate::config::settings::AppSettings; use crate::environments::{EnvironmentIndex, EnvironmentKeys}; use crate::error::{EdgeProxyError, Result}; +use crate::identities::IdentityStore; use crate::models::{ - APIFeatureState, IdentityResponse, IdentityWithTraits, ProxyConfigEnvironment, + APIFeatureState, IdentityResponse, IdentityWithTraits, ProxyConfigEnvironment, TraitModel, }; use crate::services::feature_utils::filter_out_server_key_only_flag_results; use chrono::{DateTime, Utc}; @@ -23,6 +24,7 @@ pub struct EnvironmentService { pub settings: AppSettings, pub last_updated_at: Arc>>>, environments: EnvironmentIndex, + identity_store: Option>, } impl EnvironmentService { @@ -41,6 +43,7 @@ impl EnvironmentService { settings, last_updated_at: Arc::new(RwLock::new(None)), environments, + identity_store: None, } } @@ -50,6 +53,11 @@ impl EnvironmentService { service } + pub fn with_identity_store(mut self, identity_store: Arc) -> Self { + self.identity_store = Some(identity_store); + self + } + pub async fn refresh_environment_caches(&self) -> bool { // Sync first so an environment added to the proxy config gets its // document fetched in the same pass. @@ -341,7 +349,65 @@ impl EnvironmentService { environment_key: &str, ) -> Result { let keys = self.resolve_key(environment_key)?; + match &self.identity_store { + None => { + self.evaluate_identity( + &keys, + environment_key, + &identity.identifier, + &identity.traits, + &[], + ) + .await + } + Some(identity_store) => { + self.evaluate_with_identity_store( + identity_store.as_ref(), + &keys, + environment_key, + identity, + ) + .await + } + } + } + + async fn evaluate_with_identity_store( + &self, + identity_store: &dyn IdentityStore, + keys: &EnvironmentKeys, + environment_key: &str, + identity: &IdentityWithTraits, + ) -> Result { + let environment = self + .cache + .get_environment(&keys.client_key) + .await + .ok_or_else(|| { + EdgeProxyError::ServiceUnavailable("Environment not loaded".to_string()) + })?; + let traits = identity_store.load_traits(&environment, identity).await; + let response = self + .evaluate_identity( + keys, + environment_key, + &identity.identifier, + &traits.traits, + &traits.system_traits, + ) + .await?; + identity_store.save_traits(&environment, identity, keys.key_type(environment_key)); + Ok(response) + } + async fn evaluate_identity( + &self, + keys: &EnvironmentKeys, + environment_key: &str, + identifier: &str, + traits: &[TraitModel], + system_traits: &[TraitModel], + ) -> Result { // Get pre-computed context from cache let context = self .cache @@ -351,11 +417,15 @@ impl EnvironmentService { EdgeProxyError::ServiceUnavailable("Environment not loaded".to_string()) })?; - let flagsmith_traits: Vec = - identity.traits.iter().map(Into::into).collect(); + let flagsmith_traits: Vec = traits + .iter() + .filter(|t| !system_traits.iter().any(|s| s.trait_key == t.trait_key)) + .chain(system_traits) + .map(Into::into) + .collect(); let context_with_identity = - add_identity_to_context(&context, &identity.identifier, &flagsmith_traits); + add_identity_to_context(&context, identifier, &flagsmith_traits); let evaluation_result = get_evaluation_result(&context_with_identity); @@ -373,11 +443,7 @@ impl EnvironmentService { Ok(IdentityResponse { flags: flag_results.iter().map(Into::into).collect(), - traits: identity - .traits - .iter() - .map(|t| t.to_response_json()) - .collect(), + traits: traits.iter().map(|t| t.to_response_json()).collect(), }) } diff --git a/tests/fixtures.rs b/tests/fixtures.rs index 59781a1..d6ae472 100644 --- a/tests/fixtures.rs +++ b/tests/fixtures.rs @@ -89,7 +89,7 @@ pub fn environment_1() -> serde_json::Value { { "value": "test", "operator": "EQUAL", - "property": "first_name", + "property_": "first_name", } ], "type": "ANY", diff --git a/tests/test_cors.rs b/tests/test_cors.rs index c7fa68b..c18a543 100644 --- a/tests/test_cors.rs +++ b/tests/test_cors.rs @@ -17,7 +17,7 @@ async fn server_allowing(allow_origins: &[&str]) -> TestServer { allow_origins: allow_origins.iter().map(|o| o.to_string()).collect(), ..AppSettings::new() }; - let (app, service) = create_router(settings); + let (app, service) = create_router(settings, None); service .cache .put_environment(&environment_1_api_key(), environment_1()) diff --git a/tests/test_health.rs b/tests/test_health.rs index 8016c49..d4340d2 100644 --- a/tests/test_health.rs +++ b/tests/test_health.rs @@ -7,7 +7,7 @@ use edge_proxy::routes::create_router; async fn test_liveness_check() { // Given let settings = AppSettings::new(); - let (app, _) = create_router(settings); + let (app, _) = create_router(settings, None); let server = TestServer::new(app).unwrap(); // When @@ -21,7 +21,7 @@ async fn test_liveness_check() { async fn test_health_check_returns_200_if_cache_was_updated_recently() { // Given let settings = AppSettings::new(); - let (app, service) = create_router(settings); + let (app, service) = create_router(settings, None); { let mut last_updated = service.last_updated_at.write().await; *last_updated = Some(Utc::now()); @@ -42,7 +42,7 @@ async fn test_health_check_returns_200_if_cache_was_updated_recently() { async fn test_health_check_returns_503_if_cache_was_not_updated() { // Given let settings = AppSettings::new(); - let (app, _service) = create_router(settings); + let (app, _service) = create_router(settings, None); let server = TestServer::new(app).unwrap(); // last_updated_at is None by default @@ -62,7 +62,7 @@ async fn test_health_check_returns_503_if_cache_was_not_updated() { async fn test_health_check_returns_503_if_cache_is_stale() { // Given let settings = AppSettings::new(); - let (app, service) = create_router(settings); + let (app, service) = create_router(settings, None); { let mut last_updated = service.last_updated_at.write().await; *last_updated = Some(Utc::now() - chrono::Duration::days(10)); @@ -88,7 +88,7 @@ async fn test_health_check_returns_200_if_cache_is_never_stale() { settings.health_check = HealthCheckSettings { environment_update_grace_period_seconds: None, }; - let (app, service) = create_router(settings); + let (app, service) = create_router(settings, None); let last_update_time = Utc::now() - chrono::Duration::days(10); { let mut last_updated = service.last_updated_at.write().await; diff --git a/tests/test_identity_store.rs b/tests/test_identity_store.rs new file mode 100644 index 0000000..f76b04c --- /dev/null +++ b/tests/test_identity_store.rs @@ -0,0 +1,207 @@ +mod fixtures; + +use async_trait::async_trait; +use axum_test::TestServer; +use edge_proxy::config::settings::{AppSettings, EnvironmentKeyPair}; +use edge_proxy::identities::{EnvironmentKeyType, IdentityStore, IdentityTraits}; +use edge_proxy::models::{IdentityWithTraits, TraitModel}; +use edge_proxy::routes::create_router; +use fixtures::*; +use rstest::rstest; +use serde_json::{Value, json}; +use std::sync::{Arc, Mutex}; + +#[derive(Debug, PartialEq)] +struct SaveCall { + environment_api_key: Value, + identifier: String, + key_type: EnvironmentKeyType, +} + +struct RecordingIdentityStore { + traits: IdentityTraits, + save_calls: Mutex>, +} + +impl RecordingIdentityStore { + fn new(traits: Vec, system_traits: Vec) -> Self { + Self { + traits: IdentityTraits { + traits, + system_traits, + }, + save_calls: Mutex::new(Vec::new()), + } + } +} + +#[async_trait] +impl IdentityStore for RecordingIdentityStore { + async fn load_traits( + &self, + _environment: &Value, + _identity: &IdentityWithTraits, + ) -> IdentityTraits { + self.traits.clone() + } + + fn save_traits( + &self, + environment: &Value, + identity: &IdentityWithTraits, + key_type: EnvironmentKeyType, + ) { + self.save_calls.lock().unwrap().push(SaveCall { + environment_api_key: environment["api_key"].clone(), + identifier: identity.identifier.clone(), + key_type, + }); + } +} + +fn trait_model(trait_key: &str, trait_value: &str) -> TraitModel { + TraitModel { + trait_key: trait_key.to_string(), + trait_value: json!(trait_value), + transient: false, + } +} + +async fn setup_test_server(identity_store: Option>) -> TestServer { + let settings = AppSettings { + environment_key_pairs: vec![EnvironmentKeyPair { + server_side_key: environment_1_server_key(), + client_side_key: environment_1_api_key(), + }], + api_url: "http://test".to_string(), + ..AppSettings::new() + }; + + let (app, service) = create_router(settings, identity_store); + + service + .cache + .put_environment(&environment_1_api_key(), environment_1()) + .await; + + TestServer::new(app).unwrap() +} + +async fn post_identity(server: &TestServer, environment_key: &str, traits: Value) -> Value { + let response = server + .post("/api/v1/identities/") + .add_header("X-Environment-Key", environment_key) + .json(&json!({"identifier": "identity_1", "traits": traits})) + .await; + response.assert_status_ok(); + response.json() +} + +fn feature_2_value(body: &Value) -> &Value { + &body["flags"] + .as_array() + .unwrap() + .iter() + .find(|f| f["feature"]["name"] == "feature_2") + .unwrap()["feature_state_value"] +} + +#[rstest] +#[case::segment_match("test", "segment_override")] +#[case::segment_miss("other", "2.3")] +#[tokio::test] +async fn test_post_identities_without_store_uses_request_traits( + #[case] first_name: &str, + #[case] expected_feature_2_value: &str, +) { + // Given + let server = setup_test_server(None).await; + let request_traits = json!([{"trait_key": "first_name", "trait_value": first_name}]); + + // When + let body = post_identity(&server, &environment_1_api_key(), request_traits.clone()).await; + + // Then + assert_eq!(feature_2_value(&body), expected_feature_2_value); + assert_eq!(body["traits"], request_traits); +} + +#[rstest] +#[case::trait_matches_segment(vec![trait_model("first_name", "test")], vec![], "segment_override")] +#[case::trait_misses_segment(vec![trait_model("first_name", "other")], vec![], "2.3")] +#[case::system_trait_overrides_trait( + vec![trait_model("first_name", "other")], + vec![trait_model("first_name", "test")], + "segment_override" +)] +#[tokio::test] +async fn test_post_identities_with_store_evaluates_store_traits_not_request_traits( + #[case] traits: Vec, + #[case] system_traits: Vec, + #[case] expected_feature_2_value: &str, +) { + // Given + let identity_store = Arc::new(RecordingIdentityStore::new(traits, system_traits)); + let server = setup_test_server(Some(identity_store)).await; + + // When + let body = post_identity( + &server, + &environment_1_api_key(), + json!([{"trait_key": "first_name", "trait_value": "ignored"}]), + ) + .await; + + // Then + assert_eq!(feature_2_value(&body), expected_feature_2_value); +} + +#[tokio::test] +async fn test_post_identities_with_store_returns_traits_not_request_or_system_traits() { + // Given + let identity_store = Arc::new(RecordingIdentityStore::new( + vec![trait_model("stored_key", "stored_value")], + vec![trait_model("system_key", "system_value")], + )); + let server = setup_test_server(Some(identity_store)).await; + + // When + let body = post_identity( + &server, + &environment_1_api_key(), + json!([{"trait_key": "request_key", "trait_value": "request_value"}]), + ) + .await; + + // Then + assert_eq!( + body["traits"], + json!([{"trait_key": "stored_key", "trait_value": "stored_value"}]) + ); +} + +#[rstest] +#[case::client_key(environment_1_api_key(), EnvironmentKeyType::Client)] +#[case::server_key(environment_1_server_key(), EnvironmentKeyType::Server)] +#[tokio::test] +async fn test_post_identities_with_store_saves_traits_once_with_key_type( + #[case] environment_key: String, + #[case] expected_key_type: EnvironmentKeyType, +) { + // Given + let identity_store = Arc::new(RecordingIdentityStore::new(Vec::new(), Vec::new())); + let server = setup_test_server(Some(identity_store.clone())).await; + + // When + post_identity(&server, &environment_key, json!([])).await; + + // Then + assert_eq!( + *identity_store.save_calls.lock().unwrap(), + vec![SaveCall { + environment_api_key: json!(environment_1_api_key()), + identifier: "identity_1".to_string(), + key_type: expected_key_type, + }] + ); +} diff --git a/tests/test_server.rs b/tests/test_server.rs index 4844f98..69dfabf 100644 --- a/tests/test_server.rs +++ b/tests/test_server.rs @@ -17,7 +17,7 @@ async fn setup_test_server() -> TestServer { ..AppSettings::new() }; - let (app, service) = create_router(settings); + let (app, service) = create_router(settings, None); service .cache @@ -238,7 +238,7 @@ async fn test_get_environment_document() { }], ..AppSettings::new() }; - let (app, service) = create_router(settings); + let (app, service) = create_router(settings, None); service .cache .put_environment("client_key", environment_1())