From 36f1258852a00d97a4fb27fd37d7e93d88d1f798 Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 24 Sep 2026 08:06:30 +0000 Subject: [PATCH] fix(fips): use sha256 for mcp advisory locks, oci digest and saml algorithms Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/oci/common_utils.py | 10 +---- litellm/proxy/_experimental/mcp_server/db.py | 4 +- .../management_endpoints/sso/saml_sso.py | 2 + tests/integration/_support/database.py | 18 ++++++++ tests/integration/mcp/test_mcp_management.py | 44 +++++++++++++++++++ .../integration/mcp/test_mcp_user_env_vars.py | 38 +++++++++++++++- .../mcp_server/test_mcp_env_vars.py | 18 ++++++++ .../management_endpoints/test_saml_sso.py | 15 +++++++ 8 files changed, 137 insertions(+), 12 deletions(-) diff --git a/litellm/llms/oci/common_utils.py b/litellm/llms/oci/common_utils.py index 3f703564b5a..3d973d9af7c 100644 --- a/litellm/llms/oci/common_utils.py +++ b/litellm/llms/oci/common_utils.py @@ -91,15 +91,7 @@ class OCIRequestWrapper: def sha256_base64(data: bytes) -> str: - # SHA-256 is used here to compute the x-content-sha256 header required by the - # OCI HTTP signing specification (RSA-SHA256 request signing), not for password - # or secret hashing. This is the correct and mandated algorithm for this purpose. - # See: https://docs.oracle.com/en-us/iaas/Content/API/Concepts/signingrequests.htm - # - # ``usedforsecurity=False`` declares non-security intent to static analyzers - # (CodeQL ``py/weak-sensitive-data-hashing``) — without it the request body - # gets flagged as "password-like data" via taint tracking. - digest: Final = hashlib.sha256(data, usedforsecurity=False).digest() # noqa: S324 + digest: Final = hashlib.sha256(data).digest() # noqa: S324 # OCI request-signing content digest return base64.b64encode(digest).decode() diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 70c6e6f4bf3..ab9e83e8ba5 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -593,7 +593,7 @@ def _mcp_identifier_lock_keys(*identifiers: str | None) -> tuple[int, ...]: so concurrent requests for the same pair always lock in the same order.""" return tuple( int.from_bytes( - hashlib.blake2b(f"mcp_identifier:{normalized}".encode(), digest_size=8).digest(), + hashlib.sha256(f"mcp_identifier:{normalized}".encode()).digest()[:8], "big", signed=True, ) @@ -2301,7 +2301,7 @@ async def merge_user_env_vars( """ allowed: Final = set(allowed_names) lock_key: Final = int.from_bytes( - hashlib.blake2b(f"{user_id}:{server_id}".encode(), digest_size=8).digest(), + hashlib.sha256(f"{user_id}:{server_id}".encode()).digest()[:8], "big", signed=True, ) diff --git a/litellm/proxy/management_endpoints/sso/saml_sso.py b/litellm/proxy/management_endpoints/sso/saml_sso.py index 12e1f1a03f3..a8dee1d7db2 100644 --- a/litellm/proxy/management_endpoints/sso/saml_sso.py +++ b/litellm/proxy/management_endpoints/sso/saml_sso.py @@ -201,6 +201,8 @@ class SAMLAuthHandler: "wantAssertionsSigned": SAMLAuthHandler._bool_env("SAML_WANT_ASSERTIONS_SIGNED", True), "wantMessagesSigned": SAMLAuthHandler._bool_env("SAML_WANT_MESSAGES_SIGNED", False), "authnRequestsSigned": SAMLAuthHandler._bool_env("SAML_AUTHN_REQUESTS_SIGNED", False), + "signatureAlgorithm": "http://www.w3.org/2001/04/xmldsig-more#rsa-sha256", + "digestAlgorithm": "http://www.w3.org/2001/04/xmlenc#sha256", "wantNameId": True, "requestedAuthnContext": False, "rejectUnsolicitedResponsesWithInResponseTo": False, diff --git a/tests/integration/_support/database.py b/tests/integration/_support/database.py index e7f0ebdf603..d25d66de3df 100644 --- a/tests/integration/_support/database.py +++ b/tests/integration/_support/database.py @@ -1,3 +1,4 @@ +import hashlib import os from typing import Final @@ -14,3 +15,20 @@ def read_rows( with psycopg.connect(database_url or os.environ["DATABASE_URL"], row_factory=dict_row) as connection: connection.execute("SET TRANSACTION READ ONLY") return ROWS.validate_python(connection.execute(query, parameters).fetchall()) + + +def advisory_lock_key(*parts: str) -> int: + return int.from_bytes(hashlib.sha256(":".join(parts).encode()).digest()[:8], "big", signed=True) + + +def legacy_advisory_lock_key(text: str) -> int: + return int.from_bytes(hashlib.blake2b(text.encode(), digest_size=8).digest(), "big", signed=True) + + +def advisory_waiters(lock_key: int) -> list[dict[str, JsonValue]]: + unsigned: Final = lock_key & 0xFFFFFFFFFFFFFFFF + return read_rows( + "SELECT pid FROM pg_locks WHERE locktype = 'advisory' AND objsubid = 1 " + "AND classid = %s::oid AND objid = %s::oid AND NOT granted", + (str(unsigned >> 32), str(unsigned & 0xFFFFFFFF)), + ) diff --git a/tests/integration/mcp/test_mcp_management.py b/tests/integration/mcp/test_mcp_management.py index bde18840d7d..aaeff80abe7 100644 --- a/tests/integration/mcp/test_mcp_management.py +++ b/tests/integration/mcp/test_mcp_management.py @@ -1,9 +1,13 @@ +import concurrent.futures +import os import uuid from pathlib import Path from typing import Final +import psycopg import yaml from integration._support.client import Gateway, eventually +from integration._support.database import advisory_lock_key, advisory_waiters, legacy_advisory_lock_key from integration._support.mcp import ( McpCaller, call_tool, @@ -285,3 +289,43 @@ def test_config_declared_server_behaves_like_database_server_but_is_read_only(ga assert declared_id in _servers(candidate) assert call_tool(candidate, key, declared_id, declared_names["add"], ADD).status_code == 200 assert len(tool_calls(declared_peer.drain())) == 1 and tool_calls(database_peer.drain()) == () + + +def test_create_waits_on_the_sha256_advisory_lock_for_the_identifier(gateway: Gateway) -> None: + with mcp_peer() as peer, gateway.scenario() as scenario: + name: Final = "mgmt" + uuid.uuid4().hex[:8] + lock_key: Final = advisory_lock_key(f"mcp_identifier:{name.lower()}") + with ( + concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool, + psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as holder, + ): + holder.execute("SELECT pg_advisory_lock(%s::bigint)", (lock_key,)) + pending: Final = pool.submit( + gateway.request, + "POST", + "/v1/mcp/server", + {"server_name": name, "alias": name, **peer.registration()}, + ) + eventually(lambda: advisory_waiters(lock_key), lambda rows: len(rows) == 1, seconds=20) + assert not pending.done() + holder.execute("SELECT pg_advisory_unlock(%s::bigint)", (lock_key,)) + response: Final = pending.result(timeout=30) + assert response.status_code == 201, response.text + identity: Final = str(response.json()["server_id"]) + scenario.cleanups.callback(forget_mcp, gateway, identity) + assert _servers(gateway)[identity]["server_name"] == name + + +def test_create_does_not_wait_on_the_legacy_blake2b_lock_id(gateway: Gateway) -> None: + with mcp_peer() as peer, gateway.scenario() as scenario: + name: Final = "mgmt" + uuid.uuid4().hex[:8] + lock_key: Final = legacy_advisory_lock_key(f"mcp_identifier:{name.lower()}") + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as holder: + holder.execute("SELECT pg_advisory_lock(%s::bigint)", (lock_key,)) + response: Final = gateway.request( + "POST", "/v1/mcp/server", {"server_name": name, "alias": name, **peer.registration()} + ) + assert response.status_code == 201, response.text + identity: Final = str(response.json()["server_id"]) + scenario.cleanups.callback(forget_mcp, gateway, identity) + assert _servers(gateway)[identity]["server_name"] == name diff --git a/tests/integration/mcp/test_mcp_user_env_vars.py b/tests/integration/mcp/test_mcp_user_env_vars.py index d9cecaccadb..26140d318b6 100644 --- a/tests/integration/mcp/test_mcp_user_env_vars.py +++ b/tests/integration/mcp/test_mcp_user_env_vars.py @@ -1,3 +1,4 @@ +import os import signal import uuid from collections.abc import Mapping @@ -8,9 +9,10 @@ from pathlib import Path from typing import Final import httpx +import psycopg import pytest from integration._support.client import Gateway, Scenario, eventually, object_value, string_value -from integration._support.database import read_rows +from integration._support.database import advisory_lock_key, advisory_waiters, legacy_advisory_lock_key, read_rows from integration._support.mcp import McpPeer, call_tool, mcp_peer, register_mcp, tool_names from integration._support.process import owned_proxy_process from pydantic import JsonValue, TypeAdapter @@ -357,3 +359,37 @@ def test_concurrent_stores_of_different_variables_do_not_lose_an_update(gateway: with ThreadPoolExecutor(max_workers=2) as pool: for _ in range(5): race_once(pool) + + +def test_store_waits_on_the_sha256_advisory_lock_for_the_user_and_server(gateway: Gateway) -> None: + with mcp_peer() as peer, gateway.scenario() as scenario: + identity: Final = register_user_var_server(scenario, peer, TOKEN) + user: Final = scenario.user() + key: Final = scenario.key(user_id=user, object_permission=grants(identity)) + wait_for_tools(gateway, key, identity) + lock_key: Final = advisory_lock_key(user, identity) + with ( + ThreadPoolExecutor(max_workers=1) as pool, + psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as holder, + ): + holder.execute("SELECT pg_advisory_lock(%s::bigint)", (lock_key,)) + pending: Final = pool.submit(store, gateway, key, identity, {TOKEN: "held-value"}) + eventually(lambda: advisory_waiters(lock_key), lambda rows: len(rows) == 1, seconds=20) + assert not pending.done() + holder.execute("SELECT pg_advisory_unlock(%s::bigint)", (lock_key,)) + response: Final = pending.result(timeout=30) + assert response.status_code == 200, response.text + assert set_names(env_status(gateway, key, identity)) == {TOKEN: True} + + +def test_store_does_not_wait_on_the_legacy_blake2b_lock_id(gateway: Gateway) -> None: + with mcp_peer() as peer, gateway.scenario() as scenario: + identity: Final = register_user_var_server(scenario, peer, TOKEN) + user: Final = scenario.user() + key: Final = scenario.key(user_id=user, object_permission=grants(identity)) + lock_key: Final = legacy_advisory_lock_key(f"{user}:{identity}") + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as holder: + holder.execute("SELECT pg_advisory_lock(%s::bigint)", (lock_key,)) + response: Final = store(gateway, key, identity, {TOKEN: "free-value"}) + assert response.status_code == 200, response.text + assert set_names(env_status(gateway, key, identity)) == {TOKEN: True} diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py index 93b894f7645..5c17b9665f4 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py @@ -6,6 +6,8 @@ connection. The DB-backed per-user flow is exercised in higher-level tests in tests/mcp_tests. """ +import hashlib + import pytest # Look up these names lazily on every access. Tests in this directory call @@ -16,6 +18,7 @@ import pytest # stop matching the new class. Accessing the attribute through the module # always picks up the current version. import litellm.proxy._experimental.mcp_server.utils as _mcp_utils +from litellm.proxy._experimental.mcp_server.db import _mcp_identifier_lock_keys, merge_user_env_vars def _u(name: str): @@ -1721,3 +1724,18 @@ async def test_missing_user_env_vars_error_renders_in_mcp_call_tool(): assert "CorporateDB" in text assert "CORP_USERNAME" in text assert "fill_env_vars=srv-99" in text + + +@pytest.mark.asyncio +async def test_merge_user_env_vars_locks_the_sha256_key_for_user_and_server(env_vars_salt_key): + prisma = _transactional_env_vars_prisma() + await merge_user_env_vars(prisma, "alice", "srv-1", {"TOKEN": "x"}, allowed_names=["TOKEN"]) + + expected = int.from_bytes(hashlib.sha256(b"alice:srv-1").digest()[:8], "big", signed=True) + assert list(prisma.db._store.locks) == [expected] + + +def test_mcp_identifier_lock_keys_use_sha256_of_the_lowercased_identifier(): + keys = _mcp_identifier_lock_keys("Name", "name", None) + + assert keys == (int.from_bytes(hashlib.sha256(b"mcp_identifier:name").digest()[:8], "big", signed=True),) diff --git a/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py b/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py index 635e6332958..49599ec67c1 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py @@ -22,6 +22,7 @@ from cryptography import x509 from cryptography.hazmat.primitives import hashes, serialization from cryptography.hazmat.primitives.asymmetric import rsa from cryptography.x509.oid import NameOID +from onelogin.saml2.settings import OneLogin_Saml2_Settings from onelogin.saml2.utils import OneLogin_Saml2_Utils from starlette.datastructures import URL @@ -749,3 +750,17 @@ class TestSAMLAuthnCookieSecureFlag: redirect = await SAMLAuthHandler.build_login_redirect(request, cache) cookie = redirect.headers["set-cookie"] assert "Secure" not in cookie + + +@pytest.mark.asyncio +async def test_effective_sp_settings_use_sha256_signature_and_digest(saml_env): + cache = DualCache() + idp_settings = await SAMLAuthHandler._load_idp_settings(cache) + settings = SAMLAuthHandler._build_settings(_fake_request(), idp_settings) + + security = settings["security"] + assert security["signatureAlgorithm"] == "http://www.w3.org/2001/04/xmldsig-more#rsa-sha256" + assert security["digestAlgorithm"] == "http://www.w3.org/2001/04/xmlenc#sha256" + effective = OneLogin_Saml2_Settings(settings, sp_validation_only=True).get_security_data() + assert effective["signatureAlgorithm"] == security["signatureAlgorithm"] + assert effective["digestAlgorithm"] == security["digestAlgorithm"]