fix: sanitize environment document cache payload to prevent internal data leak

Author: SahilJat

api/environments/models.py

@@ -4,6 +4,7 @@
 from copy import deepcopy
 from typing import TYPE_CHECKING, Literal
 
+from flagsmith_schemas.api import V1EnvironmentDocumentResponse
 from common.core.utils import using_database_replica
 from django.conf import settings
 from django.contrib.contenttypes.fields import GenericRelation
@@ -359,12 +360,31 @@ def write_environment_documents(
             settings.CACHE_ENVIRONMENT_DOCUMENT_MODE
             == EnvironmentDocumentCacheMode.PERSISTENT
         ):
-            environment_document_cache.set_many(
-                {
-                    e.api_key: map_environment_to_environment_document(e)
-                    for e in environments
+            cache_payload = {}
+            for e in environments:
+                fat_doc = map_environment_to_environment_document(e)
+                
+                #  Safely check for Pydantic V2's model_dump, with a fallback
+                if hasattr(fat_doc, "model_dump"):
+                    doc_dict = fat_doc.model_dump()
+                elif hasattr(fat_doc, "dict"): 
+                    # Just in case Flagsmith has a legacy custom class here
+                    doc_dict = fat_doc.dict()
+                else:
+                    doc_dict = fat_doc
+                
+                # Manually extract ONLY the keys allowed by the V1 TypedDict
+                clean_doc: V1EnvironmentDocumentResponse = {
+                    "api_key": doc_dict.get("api_key"),
+                    "feature_states": doc_dict.get("feature_states", []),
+                    "identity_overrides": doc_dict.get("identity_overrides", []),
+                    "name": doc_dict.get("name"),
+                    "project": doc_dict.get("project"),
                 }
-            )
+                cache_payload[e.api_key] = clean_doc
+
+            # Save the clean payload to Redis
+            environment_document_cache.set_many(cache_payload)
 
     def get_feature_state(
         self,

api/tests/unit/environments/test_unit_environments_models.py

@@ -20,6 +20,8 @@
 from core.constants import STRING
 from core.request_origin import RequestOrigin
 from environments.identities.models import Identity
+from environments.enums import EnvironmentDocumentCacheMode
+from flagsmith_schemas.api import V1EnvironmentDocumentResponse
 from environments.metrics import CACHE_HIT, CACHE_MISS
 from environments.models import (
     Environment,
@@ -1267,7 +1269,7 @@ def test_environment_clone_from_non_versioned_environment_with_use_v2_feature_ve
     # When
     new_environment = environment.clone(name="new-environment")
 
-    # Then
+ # Then
     assert new_environment.use_v2_feature_versioning
 
     # we only expect a single environment feature version as we are essentially
@@ -1280,3 +1282,34 @@ def test_environment_clone_from_non_versioned_environment_with_use_v2_feature_ve
     latest_feature_states = get_environment_flags_queryset(new_environment)
     assert latest_feature_states.count() == 2
     assert {fs.environment_feature_version for fs in latest_feature_states} == {efv}
+
+@mock.patch("environments.models.environment_document_cache")
+def test_write_environment_documents_strips_internal_fields_from_cache(
+    mock_document_cache: MagicMock, 
+    environment: Environment, 
+    settings: typing.Any
+) -> None:
+    # 1. SETUP: Force the persistent caching mode
+    settings.CACHE_ENVIRONMENT_DOCUMENT_MODE = EnvironmentDocumentCacheMode.PERSISTENT
+    
+    # 2. ACTION: Trigger the document builder
+    Environment.write_environment_documents(environment_id=environment.id)
+
+    # 3. INTERCEPT: Grab the exact dictionary the code tried to save to Redis
+    # call_args[0][0] gets the dictionary passed into set_many()
+    cache_payload = mock_document_cache.set_many.call_args[0][0]
+    cached_document = cache_payload[environment.api_key]
+
+    # 4. ASSERT: The internal/dangerous fields MUST NOT be in the payload
+    assert "compress_dynamo_documents" not in cached_document
+    assert "use_v2_feature_versioning" not in cached_document
+
+    # 5. SCHEMA VALIDATION: Prove the dictionary perfectly matches the TypedDict 
+    # required by the API. If this raises an error, the fix failed.
+    clean_doc: V1EnvironmentDocumentResponse = {
+        "api_key": cached_document["api_key"],
+        "feature_states": cached_document["feature_states"],
+        "identity_overrides": cached_document["identity_overrides"],
+        "name": cached_document["name"],
+        "project": cached_document["project"],
+    }