From 527d3b1c26bf8cd11349e270cdb909621bdcb52c Mon Sep 17 00:00:00 2001 From: Alphaxiaoteng <230277249+Alphaxiaoteng@users.noreply.github.com> Date: Wed, 30 Sep 2026 01:27:27 +0800 Subject: [PATCH] fix(backend): validate inputs and add resilience guards in database mcp api key --- backend/database/mcp_api_key.py | 66 +++++++-- .../tests/unit/test_mcp_api_key_resilience.py | 126 ++++++++++++++++++ 2 files changed, 180 insertions(+), 12 deletions(-) create mode 100644 backend/tests/unit/test_mcp_api_key_resilience.py diff --git a/backend/database/mcp_api_key.py b/backend/database/mcp_api_key.py index c5f82b673e3..1e159aae45d 100644 --- a/backend/database/mcp_api_key.py +++ b/backend/database/mcp_api_key.py @@ -156,6 +156,32 @@ def _valid_cached_auth_context(cached_data: Dict[str, Any]) -> bool: ) +def _clean_user_id(user_id: Any) -> str: + """Validate and clean user identifier to prevent path traversal and empty allocations.""" + if not isinstance(user_id, str): + raise ApiKeyValidationError("user_id must be a string") + cleaned = user_id.strip() + if not cleaned: + raise ApiKeyValidationError("user_id must not be empty") + if "/" in cleaned or "\\" in cleaned or ".." in cleaned or len(cleaned) > 128: + raise ApiKeyValidationError("user_id contains invalid characters or exceeds maximum length") + return cleaned + + +def _clean_key_name(name: Any) -> str: + """Validate and normalize API key name, enforcing bounds and rejecting raw token leakage.""" + if not isinstance(name, str): + raise ApiKeyValidationError("API key name must be a string") + cleaned = " ".join(name.split()).strip() + if not cleaned: + raise ApiKeyValidationError("API key name must not be empty") + if len(cleaned) > 255: + raise ApiKeyValidationError("API key name must not exceed 255 characters") + if contains_raw_api_key(cleaned): + raise ApiKeyValidationError("API key name must not contain a raw API key") + return cleaned + + def create_mcp_key( user_id: str, name: str, @@ -166,8 +192,8 @@ def create_mcp_key( Creates a new MCP API key for a user. Returns the raw key and the key's metadata. """ - if contains_raw_api_key(name): - raise ApiKeyValidationError("API key name must not contain a raw API key") + clean_uid = _clean_user_id(user_id) + clean_name = _clean_key_name(name) if app_id is None: resolved_app_id = MCP_DEFAULT_APP_ID else: @@ -182,8 +208,8 @@ def create_mcp_key( api_key_doc = { "id": key_id, - "user_id": user_id, - "name": name, + "user_id": clean_uid, + "name": clean_name, "hashed_key": hashed_key, "key_prefix": key_prefix, "created_at": now, @@ -192,11 +218,11 @@ def create_mcp_key( "scopes": resolved_scopes, } firestore_client.collection("mcp_api_keys").document(key_id).set(api_key_doc) - _seed_mcp_memory_grant(user_id, key_id, resolved_app_id, firestore_client=firestore_client) + _seed_mcp_memory_grant(clean_uid, key_id, resolved_app_id, firestore_client=firestore_client) api_key_data = McpApiKey( id=key_id, - name=name, + name=clean_name, key_prefix=key_prefix, created_at=now, last_used_at=None, @@ -212,7 +238,12 @@ def get_mcp_keys_for_user_with_repair_info( """ Retrieves all MCP API keys and bounded metadata-repair reasons for a user. """ - keys_ref = _db().collection("mcp_api_keys").where("user_id", "==", user_id) + if not isinstance(user_id, str): + return [], frozenset() + clean_uid = user_id.strip() + if not clean_uid or "/" in clean_uid or "\\" in clean_uid or ".." in clean_uid or len(clean_uid) > 128: + return [], frozenset() + keys_ref = _db().collection("mcp_api_keys").where("user_id", "==", clean_uid) docs = keys_ref.stream() keys: list[McpApiKey] = [] repairs: set[ApiKeyMetadataRepair] = set() @@ -254,13 +285,24 @@ def delete_mcp_key(user_id: str, key_id: str) -> None: """ Deletes an MCP API key. """ + if not isinstance(user_id, str) or not isinstance(key_id, str): + return + clean_uid = user_id.strip() + clean_key_id = key_id.strip() + if not clean_uid or not clean_key_id: + return + if "/" in clean_uid or "\\" in clean_uid or ".." in clean_uid or len(clean_uid) > 128: + return + if "/" in clean_key_id or "\\" in clean_key_id or ".." in clean_key_id or len(clean_key_id) > 128: + return + firestore_client = _db() - key_ref = firestore_client.collection("mcp_api_keys").document(key_id) + key_ref = firestore_client.collection("mcp_api_keys").document(clean_key_id) key_doc = key_ref.get() if key_doc.exists: raw: object = key_doc.to_dict() key_data: Dict[str, Any] = cast(Dict[str, Any], raw) if isinstance(raw, dict) else {} - if key_data.get("user_id") == user_id: + if key_data.get("user_id") == clean_uid: hashed_key = key_data.get("hashed_key") if not is_valid_api_key_hash(hashed_key): raise ApiKeyRevocationUnavailableError("MCP API key credential metadata is invalid") @@ -271,8 +313,8 @@ def delete_mcp_key(user_id: str, key_id: str) -> None: if cache_deleted is not True: raise ApiKeyRevocationUnavailableError("MCP API key cache invalidation was not confirmed") _delete_mcp_memory_grant( - user_id, - key_id, + clean_uid, + clean_key_id, normalize_api_key_app_id(key_data.get("app_id"), default=MCP_DEFAULT_APP_ID), firestore_client=firestore_client, ) @@ -298,7 +340,7 @@ def get_api_key_auth_result(api_key: str) -> ApiKeyAuthLookupResult: layer; repair them lazily on successful authentication so existing agents keep working without regenerating keys. """ - if not api_key.startswith("omi_mcp_"): + if not isinstance(api_key, str) or not api_key.startswith("omi_mcp_"): return ApiKeyAuthLookupResult(context=None) secret_part = api_key.replace("omi_mcp_", "", 1) hashed_key = hash_api_key(secret_part) diff --git a/backend/tests/unit/test_mcp_api_key_resilience.py b/backend/tests/unit/test_mcp_api_key_resilience.py new file mode 100644 index 00000000000..a5e5cf1e173 --- /dev/null +++ b/backend/tests/unit/test_mcp_api_key_resilience.py @@ -0,0 +1,126 @@ +from unittest.mock import MagicMock, patch +import pytest + +from database.api_key_metadata import ApiKeyValidationError +from database.mcp_api_key import ( + _clean_key_name, + _clean_user_id, + create_mcp_key, + delete_mcp_key, + get_api_key_auth_result, + get_mcp_keys_for_user, + get_mcp_keys_for_user_with_repair_info, +) + + +def test_clean_user_id_valid(): + assert _clean_user_id("user_12345") == "user_12345" + assert _clean_user_id(" user_abc ") == "user_abc" + + +@pytest.mark.parametrize("invalid_uid", ["", " ", None, 12345, [], {}]) +def test_clean_user_id_empty_or_invalid_type(invalid_uid): + with pytest.raises(ApiKeyValidationError, match="user_id"): + _clean_user_id(invalid_uid) + + +@pytest.mark.parametrize( + "traversal_uid", + [ + "../etc/passwd", + "users/../../root", + "user/123", + "user\\123", + "..", + "a" * 129, + ], +) +def test_clean_user_id_path_traversal_and_bounds(traversal_uid): + with pytest.raises(ApiKeyValidationError, match="invalid characters or exceeds maximum length"): + _clean_user_id(traversal_uid) + + +def test_clean_key_name_valid(): + assert _clean_key_name("Production Key") == "Production Key" + assert _clean_key_name(" Dev Key ") == "Dev Key" + + +@pytest.mark.parametrize("invalid_name", ["", " ", None, 999]) +def test_clean_key_name_empty_or_invalid_type(invalid_name): + with pytest.raises(ApiKeyValidationError, match="API key name"): + _clean_key_name(invalid_name) + + +def test_clean_key_name_bounds(): + too_long = "x" * 256 + with pytest.raises(ApiKeyValidationError, match="must not exceed 255 characters"): + _clean_key_name(too_long) + + +def test_clean_key_name_rejects_raw_token(): + raw_token = "omi_mcp_1234567890abcdef1234567890abcdef" + with pytest.raises(ApiKeyValidationError, match="must not contain a raw API key"): + _clean_key_name(f"My key: {raw_token}") + + +def test_create_mcp_key_sanitizes_and_persists(): + mock_db = MagicMock() + mock_doc = MagicMock() + mock_db.collection.return_value.document.return_value = mock_doc + + with patch("database.mcp_api_key._db", return_value=mock_db), patch( + "database.mcp_api_key._seed_mcp_memory_grant" + ) as mock_seed: + raw_key, api_key_data = create_mcp_key( + user_id=" uid_valid_42 ", + name=" Claude Agent Key ", + ) + + assert raw_key.startswith("omi_mcp_") + assert api_key_data.name == "Claude Agent Key" + mock_doc.set.assert_called_once() + saved_doc = mock_doc.set.call_args[0][0] + assert saved_doc["user_id"] == "uid_valid_42" + assert saved_doc["name"] == "Claude Agent Key" + mock_seed.assert_called_once() + assert mock_seed.call_args[0][0] == "uid_valid_42" + + +@pytest.mark.parametrize("invalid_uid", ["", " ", None, "../admin", "bad/uid", "x" * 130]) +def test_get_mcp_keys_for_user_invalid_uid_fast_return(invalid_uid): + with patch("database.mcp_api_key._db") as mock_db: + keys, repairs = get_mcp_keys_for_user_with_repair_info(invalid_uid) + assert keys == [] + assert repairs == frozenset() + mock_db.assert_not_called() + + keys_simple = get_mcp_keys_for_user(invalid_uid) + assert keys_simple == [] + + +@pytest.mark.parametrize( + "invalid_input", + [ + ("", "key_1"), + ("uid_1", ""), + (" ", "key_1"), + ("uid_1", " "), + (None, "key_1"), + ("uid_1", None), + ("../admin", "key_1"), + ("uid_1", "../etc"), + ("uid/1", "key_1"), + ("uid_1", "key/1"), + ], +) +def test_delete_mcp_key_invalid_inputs_noop(invalid_input): + uid, key_id = invalid_input + with patch("database.mcp_api_key._db") as mock_db: + delete_mcp_key(uid, key_id) + mock_db.assert_not_called() + + +@pytest.mark.parametrize("invalid_api_key", [None, 12345, [], {}, "invalid_prefix", "omi_dev_123"]) +def test_get_api_key_auth_result_non_string_or_prefix(invalid_api_key): + result = get_api_key_auth_result(invalid_api_key) + assert result.context is None