fix: support server-side keys on /flags and /identities

Cargo.toml

@@ -27,6 +27,7 @@ flagsmith-flag-engine = "0.6"
 [dev-dependencies]
 axum-test = "18"
 mockall = "0.14"
+rstest = "0.27"
 tokio-test = "0.4"
 tracing-test = { version = "0.2", features = ["no-env-filter"] }
 wiremock = "0.6"

src/services/environment.rs

@@ -296,13 +296,11 @@ impl EnvironmentService {
         environment_key: &str,
         feature_name: Option<&str>,
     ) -> Result<Vec<APIFeatureState>> {
-        // TODO: serve server-side keys here, see
-        // https://github.com/Flagsmith/edge-proxy-rs/issues/24
-        self.resolve_key(environment_key)?;
+        let keys = self.resolve_key(environment_key)?;
 
         let context = self
             .cache
-            .get_context(environment_key)
+            .get_context(&keys.client_key)
             .await
             .ok_or_else(|| {
                 EdgeProxyError::ServiceUnavailable("Environment not loaded".to_string())
@@ -319,8 +317,8 @@ impl EnvironmentService {
             }
         }
 
-        if !environment_key.starts_with("ser.") {
-            if let Some(document) = self.cache.get_environment(environment_key).await {
+        if environment_key == keys.client_key {
+            if let Some(document) = self.cache.get_environment(&keys.client_key).await {
                 let server_key_only_ids = Self::extract_server_key_only_ids(&document);
                 flag_results =
                     filter_out_server_key_only_flag_results(flag_results, &server_key_only_ids);
@@ -342,12 +340,12 @@ impl EnvironmentService {
         identity: &IdentityWithTraits,
         environment_key: &str,
     ) -> Result<IdentityResponse> {
-        self.resolve_key(environment_key)?;
+        let keys = self.resolve_key(environment_key)?;
 
         // Get pre-computed context from cache
         let context = self
             .cache
-            .get_context(environment_key)
+            .get_context(&keys.client_key)
             .await
             .ok_or_else(|| {
                 EdgeProxyError::ServiceUnavailable("Environment not loaded".to_string())
@@ -364,9 +362,9 @@ impl EnvironmentService {
         let mut flag_results: Vec<FlagResult> = evaluation_result.flags.into_values().collect();
 
         // Filter out server-key-only features if using client key
-        if !environment_key.starts_with("ser.") {
+        if environment_key == keys.client_key {
             // Get server_key_only_feature_ids from cached document
-            if let Some(document) = self.cache.get_environment(environment_key).await {
+            if let Some(document) = self.cache.get_environment(&keys.client_key).await {
                 let server_key_only_ids = Self::extract_server_key_only_ids(&document);
                 flag_results =
                     filter_out_server_key_only_flag_results(flag_results, &server_key_only_ids);

tests/fixtures.rs

@@ -4,6 +4,10 @@ pub fn environment_1_api_key() -> String {
     "environment_1_api_key".to_string()
 }
 
+pub fn environment_1_server_key() -> String {
+    "ser.test_key".to_string()
+}
+
 pub fn environment_1() -> serde_json::Value {
     json!({
         "updated_at": "1969-07-20T20:17:40Z",

tests/test_cors.rs

@@ -11,7 +11,7 @@ const ORIGIN: &str = "https://app.example.com";
 async fn server_allowing(allow_origins: &[&str]) -> TestServer {
     let settings = AppSettings {
         environment_key_pairs: vec![EnvironmentKeyPair {
-            server_side_key: "ser.test_key".to_string(),
+            server_side_key: environment_1_server_key(),
             client_side_key: environment_1_api_key(),
         }],
         allow_origins: allow_origins.iter().map(|o| o.to_string()).collect(),