From eb09aea6cae0e33cfbde2ee2a735309633b8332a Mon Sep 17 00:00:00 2001 From: Sabiha Khan Date: Mon, 29 Jun 2026 19:55:50 +0530 Subject: [PATCH] fix: masked secret saves --- api/routes/organization.py | 2 +- api/routes/user.py | 2 +- api/routes/workflow.py | 7 +- .../configuration/ai_model_configuration.py | 23 +++-- api/services/configuration/masking.py | 38 ++++++++- api/services/configuration/merge.py | 2 +- api/tests/test_masked_key_rejection.py | 83 ++++++++++++++++++- api/tests/test_resolve_effective_config.py | 5 +- 8 files changed, 143 insertions(+), 19 deletions(-) diff --git a/api/routes/organization.py b/api/routes/organization.py index cc6f02ac..c8afd639 100644 --- a/api/routes/organization.py +++ b/api/routes/organization.py @@ -326,7 +326,7 @@ async def save_model_configuration_v2( existing = await get_organization_ai_model_configuration_v2(organization_id) configuration = merge_ai_model_configuration_v2_secrets(request, existing) try: - check_for_masked_keys_in_ai_model_configuration_v2(configuration) + check_for_masked_keys_in_ai_model_configuration_v2(configuration, existing) effective = compile_ai_model_configuration_v2(configuration) await UserConfigurationValidator().validate( effective, diff --git a/api/routes/user.py b/api/routes/user.py index 4a2caa4b..5e956432 100644 --- a/api/routes/user.py +++ b/api/routes/user.py @@ -155,7 +155,7 @@ async def update_user_configurations( raise HTTPException(status_code=422, detail=str(e)) try: - check_for_masked_keys(user_configurations) + check_for_masked_keys(user_configurations, existing_config) except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) diff --git a/api/routes/workflow.py b/api/routes/workflow.py index 50b17cf7..191f0438 100644 --- a/api/routes/workflow.py +++ b/api/routes/workflow.py @@ -1078,6 +1078,7 @@ async def update_workflow( incoming_v2_override, existing_v2_override_config, ) + mask_check_existing = existing_v2_override_config if existing_v2_override_config is None: resolved_config = await get_resolved_ai_model_configuration( user_id=user.id, @@ -1087,7 +1088,11 @@ async def update_workflow( v2_override, resolved_config.organization_configuration, ) - check_for_masked_keys_in_ai_model_configuration_v2(v2_override) + mask_check_existing = resolved_config.organization_configuration + check_for_masked_keys_in_ai_model_configuration_v2( + v2_override, + mask_check_existing, + ) effective = compile_ai_model_configuration_v2(v2_override) await UserConfigurationValidator().validate( effective, diff --git a/api/services/configuration/ai_model_configuration.py b/api/services/configuration/ai_model_configuration.py index 3aabdae6..c4dbba8e 100644 --- a/api/services/configuration/ai_model_configuration.py +++ b/api/services/configuration/ai_model_configuration.py @@ -244,7 +244,10 @@ def merge_ai_model_configuration_v2_secrets( existing_dograh = existing_dict.get("dograh") or {} incoming_key = incoming_dograh.get("api_key") existing_key = existing_dograh.get("api_key") - if incoming_key and existing_key and contains_masked_key(incoming_key): + if incoming_key and existing_key and contains_masked_key( + incoming_key, + existing_key, + ): incoming_dograh["api_key"] = resolve_masked_api_keys( incoming_key, existing_key, @@ -258,9 +261,13 @@ def merge_ai_model_configuration_v2_secrets( def check_for_masked_keys_in_ai_model_configuration_v2( configuration: OrganizationAIModelConfigurationV2, + existing: OrganizationAIModelConfigurationV2 | None = None, ) -> None: data = configuration.model_dump(mode="json", exclude_none=True) - _raise_if_masked_secret(data) + existing_data = ( + existing.model_dump(mode="json", exclude_none=True) if existing else None + ) + _raise_if_masked_secret(data, existing_data) def mask_ai_model_configuration_v2( @@ -381,22 +388,26 @@ def _merge_service_secret_fields(incoming: dict, existing: dict): existing_secret = existing[secret_field] if incoming_secret is None: incoming[secret_field] = existing_secret - elif contains_masked_key(incoming_secret): + elif contains_masked_key(incoming_secret, existing_secret): incoming[secret_field] = resolve_masked_api_keys( incoming_secret, existing_secret, ) -def _raise_if_masked_secret(value): +def _raise_if_masked_secret(value, existing=None): if isinstance(value, dict): for key, nested in value.items(): - if key in SERVICE_SECRET_FIELDS and contains_masked_key(nested): + existing_nested = existing.get(key) if isinstance(existing, dict) else None + if key in SERVICE_SECRET_FIELDS and contains_masked_key( + nested, + existing_nested, + ): raise ValueError( f"The {key} appears to be masked. Please provide the actual " "value, not the masked value." ) - _raise_if_masked_secret(nested) + _raise_if_masked_secret(nested, existing_nested) elif isinstance(value, list): for item in value: _raise_if_masked_secret(item) diff --git a/api/services/configuration/masking.py b/api/services/configuration/masking.py index bbcd983a..c0c23277 100644 --- a/api/services/configuration/masking.py +++ b/api/services/configuration/masking.py @@ -23,12 +23,23 @@ SERVICE_SECRET_FIELDS = ("api_key", "credentials", "aws_access_key", "aws_secret MODEL_OVERRIDE_FIELDS = ("llm", "tts", "stt", "realtime") -def contains_masked_key(value: str | list[str] | None) -> bool: +def contains_masked_key( + value: str | list[str] | None, + existing: str | list[str] | None = None, +) -> bool: """Return True if *value* looks like a masked placeholder.""" if value is None: return False keys = value if isinstance(value, list) else [value] - return any(MASK_MARKER in k or _matches_mask_shape(k) for k in keys) + if existing is None: + return any(MASK_MARKER in k for k in keys) + + existing_keys = existing if isinstance(existing, list) else [existing] + return any( + MASK_MARKER in key + or any(_matches_existing_mask(key, real_key) for real_key in existing_keys) + for key in keys + ) def _matches_mask_shape(key: str) -> bool: @@ -43,20 +54,39 @@ def _matches_mask_shape(key: str) -> bool: return key[:masked_prefix_length] == MASK_CHAR * masked_prefix_length -def check_for_masked_keys(config: "EffectiveAIModelConfiguration") -> None: +def _matches_existing_mask(key: str, existing: str | None) -> bool: + return bool(existing) and _matches_mask_shape(key) and is_mask_of(key, existing) + + +def check_for_masked_keys( + config: "EffectiveAIModelConfiguration", + existing: "EffectiveAIModelConfiguration | None" = None, +) -> None: """Raise ValueError if any service in *config* still has a masked secret.""" for field in ("llm", "tts", "stt", "embeddings", "realtime"): service = getattr(config, field, None) if service is None: continue + existing_service = getattr(existing, field, None) if existing else None for secret_field in SERVICE_SECRET_FIELDS: if not hasattr(service, secret_field): continue if secret_field == "api_key" and hasattr(service, "get_all_api_keys"): secret_value = service.get_all_api_keys() + existing_secret_value = ( + existing_service.get_all_api_keys() + if existing_service + and hasattr(existing_service, "get_all_api_keys") + else None + ) else: secret_value = getattr(service, secret_field, None) - if contains_masked_key(secret_value): + existing_secret_value = ( + getattr(existing_service, secret_field, None) + if existing_service and hasattr(existing_service, secret_field) + else None + ) + if contains_masked_key(secret_value, existing_secret_value): raise ValueError( f"The {field} {secret_field} appears to be masked. " "Please provide the actual value, not the masked value." diff --git a/api/services/configuration/merge.py b/api/services/configuration/merge.py index 3100fa45..c73b62f7 100644 --- a/api/services/configuration/merge.py +++ b/api/services/configuration/merge.py @@ -50,7 +50,7 @@ def _merge_service_secret_fields( incoming_secret = incoming_cfg.get(secret_field) existing_secret = existing_cfg[secret_field] if incoming_secret is not None: - if contains_masked_key(incoming_secret): + if contains_masked_key(incoming_secret, existing_secret): incoming_cfg[secret_field] = ( existing_secret if masked_value_preserves_full_secret diff --git a/api/tests/test_masked_key_rejection.py b/api/tests/test_masked_key_rejection.py index b1c8fe5a..57cc91a0 100644 --- a/api/tests/test_masked_key_rejection.py +++ b/api/tests/test_masked_key_rejection.py @@ -45,16 +45,16 @@ def _existing_openai_config(): class TestMaskedKeyRejection: def test_detects_short_masked_api_keys(self): """Short masked API keys should be rejected even without the legacy marker.""" - assert contains_masked_key(mask_key("EMPTY")) - assert contains_masked_key(mask_key("X")) - assert contains_masked_key(mask_key("mykey")) + assert contains_masked_key(mask_key("EMPTY"), "EMPTY") + assert contains_masked_key(mask_key("X"), "X") + assert contains_masked_key(mask_key("mykey"), "mykey") def test_detects_masked_api_key_with_mask_char_in_visible_suffix(self): """The visible suffix can contain MASK_CHAR and still be masked.""" masked = mask_key("vsk*v") assert masked == "*sk*v" - assert contains_masked_key(masked) + assert contains_masked_key(masked, "vsk*v") def test_allows_unmasked_short_api_keys(self): """Unmasked short keys should not be rejected as masked placeholders.""" @@ -62,6 +62,8 @@ class TestMaskedKeyRejection: assert not contains_masked_key("X") assert not contains_masked_key("mykey") assert not contains_masked_key("vsk*v") + assert not contains_masked_key("*rk5f") + assert not contains_masked_key("*rk5f", "different-existing-key") def test_rejects_masked_api_key_on_provider_change(self): """Changing provider with a masked API key should return 400.""" @@ -94,6 +96,42 @@ class TestMaskedKeyRejection: assert response.status_code == 400 assert "masked" in response.json()["detail"].lower() + def test_rejects_short_masked_api_key_on_provider_change(self): + """Changing provider with a short masked API key should return 400.""" + app = _make_test_app() + client = TestClient(app) + existing = EffectiveAIModelConfiguration( + llm=OpenAILLMService( + provider="openai", + api_key="EMPTY", + model="gpt-4.1", + ) + ) + + with ( + patch("api.routes.user.db_client") as mock_db, + patch("api.routes.user.UserConfigurationValidator") as mock_validator, + ): + mock_db.get_user_configurations = AsyncMock(return_value=existing) + mock_db.update_user_configuration = AsyncMock( + side_effect=lambda uid, cfg: cfg + ) + mock_validator.return_value.validate = AsyncMock() + + response = client.put( + "/user/configurations/user", + json={ + "llm": { + "provider": "google", + "api_key": mask_key("EMPTY"), + "model": "gemini-2.5-flash", + } + }, + ) + + assert response.status_code == 400 + assert "masked" in response.json()["detail"].lower() + def test_rejects_masked_api_key_in_list(self): """A list of API keys containing a masked key should return 400.""" app = _make_test_app() @@ -162,6 +200,43 @@ class TestMaskedKeyRejection: assert response.status_code == 200 + def test_allows_new_star_prefixed_api_key(self): + """A new real key may have the same shape as an unrelated masked key.""" + app = _make_test_app() + client = TestClient(app) + + new_key = "*rk5f" + updated = EffectiveAIModelConfiguration( + llm=GoogleLLMService( + provider="google", + api_key=new_key, + model="gemini-2.5-flash", + ) + ) + + with ( + patch("api.routes.user.db_client") as mock_db, + patch("api.routes.user.UserConfigurationValidator") as mock_validator, + ): + mock_db.get_user_configurations = AsyncMock( + return_value=_existing_openai_config() + ) + mock_db.update_user_configuration = AsyncMock(return_value=updated) + mock_validator.return_value.validate = AsyncMock() + + response = client.put( + "/user/configurations/user", + json={ + "llm": { + "provider": "google", + "api_key": new_key, + "model": "gemini-2.5-flash", + } + }, + ) + + assert response.status_code == 200 + def test_allows_same_provider_with_masked_key(self): """Same provider with masked key should succeed (merge resolves it).""" app = _make_test_app() diff --git a/api/tests/test_resolve_effective_config.py b/api/tests/test_resolve_effective_config.py index 1b9ad8c6..09875424 100644 --- a/api/tests/test_resolve_effective_config.py +++ b/api/tests/test_resolve_effective_config.py @@ -538,7 +538,10 @@ class TestWorkflowConfigurationSecrets: masked = mask_workflow_configurations(configs) assert masked["model_overrides"]["llm"]["api_key"] != "sk-real-llm-key" - assert contains_masked_key(masked["model_overrides"]["llm"]["api_key"]) + assert contains_masked_key( + masked["model_overrides"]["llm"]["api_key"], + "sk-real-llm-key", + ) assert masked["model_overrides"]["llm"]["api_key"].endswith("-key") assert masked["model_overrides"]["tts"]["api_key"] != "el-real-tts-key" assert masked["ambient_noise_configuration"] == {"enabled": True}