diff --git a/litellm/proxy/common_utils/fips.py b/litellm/proxy/common_utils/fips.py new file mode 100644 index 00000000000..8c70b7d7825 --- /dev/null +++ b/litellm/proxy/common_utils/fips.py @@ -0,0 +1,125 @@ +import os +from typing import Final + +from litellm._logging import verbose_proxy_logger +from litellm.repositories.fips_repository import FIPS_MARKER, FipsRepository +from litellm.secret_managers.main import str_to_bool + +FIPS_MODE_ENV_VAR: Final = "LITELLM_FIPS_MODE" + + +class FipsModeError(Exception): + pass + + +def is_fips_mode() -> bool: + return str_to_bool(os.getenv(FIPS_MODE_ENV_VAR)) is True + + +def openssl_enforces_fips() -> bool: + """Python's OpenSSL runs with FIPS-only defaults and itself refuses MD5.""" + try: + import _hashlib + + if _hashlib.get_fips_mode() != 1: + return False + except (ImportError, AttributeError, ValueError): + return False + try: + _hashlib.new("md5", b"", usedforsecurity=True) + except ValueError: + return True + return False + + +def cryptography_enforces_fips() -> bool: + """The cryptography library's own OpenSSL runs with FIPS-only defaults and its FIPS provider can hash SHA-256.""" + try: + from cryptography.exceptions import InternalError, UnsupportedAlgorithm + from cryptography.hazmat.bindings._rust import openssl as rust_openssl + from cryptography.hazmat.primitives import hashes + except ImportError: + return False + try: + if not rust_openssl.is_fips_enabled(): + return False + hashes.Hash(hashes.SHA256()).finalize() + except (AttributeError, InternalError, UnsupportedAlgorithm): + return False + return True + + +def fips_mode_refusal() -> str | None: + """Why the proxy must not start: LITELLM_FIPS_MODE is not a boolean, or is on but Python's or the + cryptography library's OpenSSL does not enforce FIPS. None when it may start.""" + value: Final = os.getenv(FIPS_MODE_ENV_VAR) + if value and str_to_bool(value) is None: + return f"{FIPS_MODE_ENV_VAR}={value!r} is not a boolean. Set it to true or false, or unset it." + if not is_fips_mode(): + return None + if not openssl_enforces_fips(): + return ( + f"{FIPS_MODE_ENV_VAR} is on but OpenSSL does not enforce FIPS. Run the proxy on an OpenSSL with " + f"its FIPS provider active and default_properties = fips=yes, or unset {FIPS_MODE_ENV_VAR}." + ) + if not cryptography_enforces_fips(): + return ( + f"{FIPS_MODE_ENV_VAR} is on but the cryptography library's OpenSSL has no working FIPS provider. Build it " + "against the system OpenSSL (pip install --no-binary cryptography cryptography), or unset " + f"{FIPS_MODE_ENV_VAR}." + ) + return None + + +def enforce_fips_mode() -> None: + """Raise FipsModeError with the fips_mode_refusal reason, if there is one.""" + refusal: Final = fips_mode_refusal() + if refusal is not None: + raise FipsModeError(refusal) + + +async def _is_marked_fips(database: FipsRepository) -> bool: + try: + return await database.is_marked_fips() + except Exception as error: # noqa: BLE001 # without FIPS mode, an unreadable database starts as it did before + verbose_proxy_logger.warning("Could not check whether the database was set up in FIPS mode: %s", error) + return False + + +async def fips_database_refusal(database: FipsRepository | None, *, database_is_configured: bool) -> str | None: + """Why the proxy must not start on its database: LITELLM_FIPS_MODE is off on a database set up in FIPS mode, or + on over a database that can't be checked or already holds rows written without it. None when it may start.""" + if not is_fips_mode(): + if database is not None and await _is_marked_fips(database): + return ( + f"This database was set up with {FIPS_MODE_ENV_VAR}=true, so the proxy only starts on it in FIPS mode. " + f"To stop using FIPS mode on it for good, run " + f"DELETE FROM \"LiteLLM_Config\" WHERE param_name = '{FIPS_MARKER}'." + ) + return None + if database is None: + return ( + f"{FIPS_MODE_ENV_VAR} is on but the database can't be reached to check it." + if database_is_configured + else None + ) + try: + state: Final = await database.state() + except Exception as error: # noqa: BLE001 # a value, so FIPS mode refuses a database it can't check + return f"{FIPS_MODE_ENV_VAR} is on but the database could not be checked: {error}" + if state.holds_rows and not state.marked_fips: + return ( + f"{FIPS_MODE_ENV_VAR} is on but this database already holds data written without it. " + "FIPS mode only starts on a new, empty database." + ) + return None + + +async def enforce_fips_database(database: FipsRepository | None, *, database_is_configured: bool) -> None: + """Raise FipsModeError with the fips_database_refusal reason, if there is one, then mark the database as set up + in FIPS mode when LITELLM_FIPS_MODE is on.""" + refusal: Final = await fips_database_refusal(database, database_is_configured=database_is_configured) + if refusal is not None: + raise FipsModeError(refusal) + if database is not None and is_fips_mode(): + await database.mark_fips() diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1f1ab52e70d..84fb877a9f1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -499,6 +499,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( encrypt_value_helper, ) from litellm.proxy.common_utils.error_body_call_id import JSON_OBJECT, error_body_call_id, with_call_id +from litellm.proxy.common_utils.fips import enforce_fips_database, enforce_fips_mode from litellm.proxy.common_utils.healthy_model_filter import ( get_hidden_unhealthy_model_names, is_healthy_only_listing_default, @@ -991,6 +992,7 @@ from litellm.proxy.utils import ( # noqa: F401, RUF100 # legacy module exports from litellm.proxy.video_endpoints.endpoints import router as video_router from litellm.repositories.base_repository import SupportsModelDump from litellm.repositories.credentials_repository import CredentialsRepository +from litellm.repositories.fips_repository import FipsRepository from litellm.repositories.prisma_protocols import TableActions from litellm.responses.streaming_iterator import drain_logging_tasks from litellm.router import ( @@ -1605,6 +1607,8 @@ async def _proxy_lifespan(app: FastAPI, resources: AsyncExitStack) -> AsyncGener if decoded_worker_config is not None: await initialize_from_worker_config(decoded_worker_config) + enforce_fips_mode() + async with AsyncExitStack() as boot_database: enforce_master_key_boot_verdict( await with_stored_secrets_counted( @@ -1637,6 +1641,11 @@ async def _proxy_lifespan(app: FastAPI, resources: AsyncExitStack) -> AsyncGener database: Final = existing_database or await _own_database(resources, prisma_client) + await enforce_fips_database( + None if prisma_client is None else FipsRepository(prisma_client), + database_is_configured=get_secret("DATABASE_URL", None) is not None, + ) + await migrate_if_requested( environ=os.environ, master_key=master_key, diff --git a/litellm/repositories/fips_repository.py b/litellm/repositories/fips_repository.py new file mode 100644 index 00000000000..23a756e44bf --- /dev/null +++ b/litellm/repositories/fips_repository.py @@ -0,0 +1,65 @@ +"""FIPS repository for the LiteLLM_Config row that marks a database as set up in FIPS mode.""" + +from dataclasses import dataclass +from typing import Final, Protocol, cast # noqa: TID251 # Prisma delegates query_raw and execute_raw dynamically + +from litellm.proxy.db.create_views import SupportsRawQueries + +FIPS_MARKER: Final = "fips_mode" +_IS_MARKED: Final = 'SELECT 1 FROM "LiteLLM_Config" WHERE param_name = $1' +_OTHER_TABLES: Final = ( + "SELECT table_name FROM information_schema.tables WHERE table_schema = ANY (current_schemas(false)) " + "AND table_type = 'BASE TABLE' AND table_name LIKE 'LiteLLM\\_%' AND table_name <> 'LiteLLM_Config'" +) +_MARK: Final = ( + "INSERT INTO \"LiteLLM_Config\" (param_name, param_value) VALUES ($1, 'true'::jsonb) " + "ON CONFLICT (param_name) DO NOTHING" +) + + +def _quoted(identifier: str) -> str: + return '"' + identifier.replace('"', '""') + '"' + + +class _WithWriter(Protocol): + @property + def writer_db(self) -> object: ... + + +@dataclass(frozen=True, slots=True) +class FipsDatabaseState: + marked_fips: bool + holds_rows: bool + + +class FipsRepository: + """Reads and writes go to the writer, since a replica may not have the marker yet.""" + + def __init__(self, prisma_client: _WithWriter) -> None: + self._prisma_client: Final = prisma_client + + @property + def _writer(self) -> SupportsRawQueries: + writer: Final = self._prisma_client.writer_db + return cast(SupportsRawQueries, writer) # cast-ok: Prisma delegates database methods dynamically + + async def is_marked_fips(self) -> bool: + return bool(await self._writer.query_raw(_IS_MARKED, FIPS_MARKER)) + + async def state(self) -> FipsDatabaseState: + """The marker and whether any other LiteLLM row exists, read in one statement so they agree.""" + writer: Final = self._writer + tables: Final = tuple(str(row["table_name"]) for row in await writer.query_raw(_OTHER_TABLES)) + any_row_in_other_tables: Final = "".join(f" OR EXISTS (SELECT 1 FROM {_quoted(table)})" for table in tables) + row: Final = ( + await writer.query_raw( + 'SELECT EXISTS (SELECT 1 FROM "LiteLLM_Config" WHERE param_name = $1) AS marked_fips, ' + f'(EXISTS (SELECT 1 FROM "LiteLLM_Config" WHERE param_name <> $1){any_row_in_other_tables}) ' + "AS holds_rows", + FIPS_MARKER, + ) + )[0] + return FipsDatabaseState(marked_fips=row["marked_fips"] is True, holds_rows=row["holds_rows"] is True) + + async def mark_fips(self) -> None: + await self._writer.execute_raw(_MARK, FIPS_MARKER) diff --git a/tests/integration/configuration/test_fips_mode_boot.py b/tests/integration/configuration/test_fips_mode_boot.py new file mode 100644 index 00000000000..11e1f8fc088 --- /dev/null +++ b/tests/integration/configuration/test_fips_mode_boot.py @@ -0,0 +1,67 @@ +import _hashlib +import json +from pathlib import Path +from typing import Final + +import pytest + +from tests.integration._support.client import Gateway +from tests.integration._support.database import scratch_database, write_rows +from tests.integration._support.process import refused_boot_log + + +def _config(directory: Path, upstream_url: str) -> Path: + config: Final = directory / "fips_boot.yaml" + config.write_text( + json.dumps( + { + "model_list": [ + { + "model_name": "fips-model", + "litellm_params": {"model": "openai/fips-model", "api_base": upstream_url, "api_key": "sk-x"}, + } + ], + } + ) + ) + return config + + +def test_fips_mode_refuses_to_start_when_openssl_does_not_enforce_fips(gateway: Gateway, tmp_path: Path) -> None: + if _hashlib.get_fips_mode() == 1: + pytest.skip("this runner's OpenSSL is in FIPS mode") + log: Final = refused_boot_log( + gateway, tmp_path, {"LITELLM_FIPS_MODE": "true"}, config=_config(tmp_path, gateway.upstream_url) + ) + assert "LITELLM_FIPS_MODE is on but OpenSSL does not enforce FIPS" in log, log + + +def test_fips_mode_refuses_to_start_on_a_value_that_is_not_a_boolean(gateway: Gateway, tmp_path: Path) -> None: + log: Final = refused_boot_log( + gateway, tmp_path, {"LITELLM_FIPS_MODE": "enforced"}, config=_config(tmp_path, gateway.upstream_url) + ) + assert "LITELLM_FIPS_MODE='enforced' is not a boolean" in log, log + + +def test_proxy_refuses_to_start_without_fips_mode_on_a_database_set_up_in_fips_mode( + gateway: Gateway, tmp_path: Path +) -> None: + with scratch_database() as database_url: + write_rows( + 'CREATE TABLE "LiteLLM_Config" (param_name TEXT NOT NULL PRIMARY KEY, param_value JSONB, ' + "last_run_at TIMESTAMP(3), reload_revision BIGINT NOT NULL DEFAULT 0)", + (), + database_url=database_url, + ) + write_rows( + "INSERT INTO \"LiteLLM_Config\" (param_name, param_value) VALUES ('fips_mode', 'true')", + (), + database_url=database_url, + ) + log: Final = refused_boot_log( + gateway, + tmp_path, + {"LITELLM_FIPS_MODE": "false", "DATABASE_URL": database_url}, + config=_config(tmp_path, gateway.upstream_url), + ) + assert "This database was set up with LITELLM_FIPS_MODE=true" in log, log diff --git a/tests/integration/database/test_fips_database.py b/tests/integration/database/test_fips_database.py new file mode 100644 index 00000000000..54238bc6876 --- /dev/null +++ b/tests/integration/database/test_fips_database.py @@ -0,0 +1,93 @@ +import asyncio +from contextlib import nullcontext +from typing import Final, LiteralString + +import pytest + +from litellm.proxy.common_utils.fips import FipsModeError, enforce_fips_database +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.repositories.fips_repository import FipsRepository +from tests.integration._support.database import read_rows, scratch_database, write_rows + +_CONFIG_TABLE: Final = ( + 'CREATE TABLE "LiteLLM_Config" (param_name TEXT PRIMARY KEY, param_value JSONB, last_run_at TIMESTAMP(3), ' + "reload_revision BIGINT NOT NULL DEFAULT 0)" +) +_KEY_TABLE: Final = 'CREATE TABLE "LiteLLM_VerificationToken" (token TEXT PRIMARY KEY)' +_MARKER: Final = "INSERT INTO \"LiteLLM_Config\" (param_name, param_value) VALUES ('fips_mode', 'true')" +_KEY: Final = "INSERT INTO \"LiteLLM_VerificationToken\" (token) VALUES ('hashed-key')" +_SETTING: Final = "INSERT INTO \"LiteLLM_Config\" (param_name, param_value) VALUES ('general_settings', '{}')" + + +def _marker_rows(database_url: str) -> int: + return len( + read_rows('SELECT 1 FROM "LiteLLM_Config" WHERE param_name = %s', ("fips_mode",), database_url=database_url) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "fips_mode, setup, refusal, marked_after", + [ + ("true", (_CONFIG_TABLE, _KEY_TABLE), None, 1), + ("true", (_CONFIG_TABLE, _KEY_TABLE, _MARKER, _KEY), None, 1), + ("true", (_CONFIG_TABLE, _KEY_TABLE, _KEY), "already holds data written without it", 0), + ("true", (_CONFIG_TABLE, _KEY_TABLE, _SETTING), "already holds data written without it", 0), + ("true", (), "could not be checked", None), + ("false", (_CONFIG_TABLE, _KEY_TABLE, _MARKER), "set up with LITELLM_FIPS_MODE=true", 1), + ("false", (_CONFIG_TABLE, _KEY_TABLE, _KEY), None, 0), + ("false", (), None, None), + ], + ids=[ + "fips-on-new-database", + "fips-on-fips-database", + "fips-on-database-with-a-key", + "fips-on-database-with-a-setting", + "fips-on-database-without-tables", + "fips-off-fips-database", + "fips-off-database-with-a-key", + "fips-off-database-without-tables", + ], +) +async def test_fips_database_check( + monkeypatch: pytest.MonkeyPatch, + fips_mode: str, + setup: tuple[LiteralString, ...], + refusal: str | None, + marked_after: int | None, +) -> None: + monkeypatch.setenv("LITELLM_FIPS_MODE", fips_mode) + with scratch_database() as database_url: + for statement in setup: + write_rows(statement, (), database_url=database_url) + monkeypatch.setenv("DATABASE_URL", database_url) + client: Final = PrismaClient(database_url, ProxyLogging(UserApiKeyCache())) + await client.connect() + try: + with pytest.raises(FipsModeError, match=refusal) if refusal else nullcontext(): + await enforce_fips_database(FipsRepository(client), database_is_configured=True) + finally: + await client.aclose() + if marked_after is not None: + assert _marker_rows(database_url) == marked_after + + +@pytest.mark.asyncio +async def test_pods_starting_together_in_fips_mode_on_a_new_database_both_start( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("LITELLM_FIPS_MODE", "true") + with scratch_database() as database_url: + for statement in (_CONFIG_TABLE, _KEY_TABLE): + write_rows(statement, (), database_url=database_url) + monkeypatch.setenv("DATABASE_URL", database_url) + pods: Final = tuple(PrismaClient(database_url, ProxyLogging(UserApiKeyCache())) for _ in range(4)) + await asyncio.gather(*(pod.connect() for pod in pods)) + try: + await asyncio.gather( + *(enforce_fips_database(FipsRepository(pod), database_is_configured=True) for pod in pods) + ) + finally: + await asyncio.gather(*(pod.aclose() for pod in pods)) + assert _marker_rows(database_url) == 1 diff --git a/tests/unit/proxy/common_utils/test_fips.py b/tests/unit/proxy/common_utils/test_fips.py new file mode 100644 index 00000000000..a48ba3138e2 --- /dev/null +++ b/tests/unit/proxy/common_utils/test_fips.py @@ -0,0 +1,195 @@ +import _hashlib +from collections.abc import Mapping, Sequence +from contextlib import nullcontext +from typing import Final +from unittest.mock import Mock + +import pytest +from cryptography.exceptions import InternalError +from cryptography.hazmat.bindings._rust import openssl as rust_openssl +from cryptography.hazmat.primitives import hashes + +from litellm.proxy.common_utils.fips import ( + FipsModeError, + cryptography_enforces_fips, + enforce_fips_database, + enforce_fips_mode, + is_fips_mode, + openssl_enforces_fips, +) +from litellm.repositories.fips_repository import FipsRepository + + +def _fake_openssl(monkeypatch: pytest.MonkeyPatch, *, fips_flag: int, md5_refused: bool) -> None: + real_new: Final = _hashlib.new + + def new(name: str, data: bytes = b"", *, usedforsecurity: bool = True) -> _hashlib.HASH: + if name == "md5" and md5_refused: + raise _hashlib.UnsupportedDigestmodError("unsupported hash type") + return real_new(name, data, usedforsecurity=usedforsecurity) + + monkeypatch.setattr(_hashlib, "get_fips_mode", lambda: fips_flag) + monkeypatch.setattr(_hashlib, "new", new) + + +@pytest.mark.parametrize("value, expected", [("true", True), (" TRUE ", True), ("false", False), (None, False)]) +def test_is_fips_mode(monkeypatch: pytest.MonkeyPatch, value: str | None, expected: bool) -> None: + if value is None: + monkeypatch.delenv("LITELLM_FIPS_MODE", raising=False) + else: + monkeypatch.setenv("LITELLM_FIPS_MODE", value) + assert is_fips_mode() is expected + + +@pytest.mark.parametrize( + "fips_flag, md5_refused, enforces", + [(1, True, True), (0, True, False), (1, False, False)], + ids=["fips", "fips-flag-off", "md5-allowed"], +) +def test_openssl_enforces_fips_needs_the_fips_flag_and_a_refused_md5( + monkeypatch: pytest.MonkeyPatch, fips_flag: int, md5_refused: bool, enforces: bool +) -> None: + _fake_openssl(monkeypatch, fips_flag=fips_flag, md5_refused=md5_refused) + assert openssl_enforces_fips() is enforces + + +def test_openssl_without_get_fips_mode_does_not_enforce_fips(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delattr(_hashlib, "get_fips_mode") + assert openssl_enforces_fips() is False + + +@pytest.mark.parametrize( + "fips_enabled, sha256_broken, enforces", + [(True, False, True), (False, False, False), (True, True, False)], + ids=["fips", "fips-flag-off", "sha256-broken"], +) +def test_cryptography_enforces_fips_needs_the_fips_flag_and_a_working_sha256( + monkeypatch: pytest.MonkeyPatch, fips_enabled: bool, sha256_broken: bool, enforces: bool +) -> None: + monkeypatch.setattr(rust_openssl, "is_fips_enabled", lambda: fips_enabled) + if sha256_broken: + monkeypatch.setattr(hashes, "Hash", Mock(side_effect=InternalError("Unknown OpenSSL error", []))) + assert cryptography_enforces_fips() is enforces + + +def test_cryptography_without_is_fips_enabled_does_not_enforce_fips(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delattr(rust_openssl, "is_fips_enabled") + assert cryptography_enforces_fips() is False + + +def test_this_runners_python_openssl_does_not_enforce_fips() -> None: + if _hashlib.get_fips_mode() == 1: + pytest.skip("this runner's Python OpenSSL is in FIPS mode") + assert openssl_enforces_fips() is False + + +def test_this_runners_cryptography_openssl_does_not_enforce_fips() -> None: + if rust_openssl.is_fips_enabled(): + pytest.skip("this runner's cryptography OpenSSL is in FIPS mode") + assert cryptography_enforces_fips() is False + + +@pytest.mark.parametrize( + "fips_mode, openssl_fips, cryptography_fips, refusal", + [ + (None, False, False, None), + ("", False, False, None), + ("false", False, False, None), + ("true", True, True, None), + ("true", False, True, "is on but OpenSSL does not enforce FIPS"), + ("true", True, False, "cryptography library's OpenSSL has no working FIPS provider"), + ("enforced", True, True, "LITELLM_FIPS_MODE='enforced' is not a boolean"), + ], + ids=["unset", "empty", "false", "fips", "openssl-not-fips", "cryptography-not-fips", "not-a-boolean"], +) +def test_enforce_fips_mode( + monkeypatch: pytest.MonkeyPatch, + fips_mode: str | None, + openssl_fips: bool, + cryptography_fips: bool, + refusal: str | None, +) -> None: + if fips_mode is None: + monkeypatch.delenv("LITELLM_FIPS_MODE", raising=False) + else: + monkeypatch.setenv("LITELLM_FIPS_MODE", fips_mode) + _fake_openssl(monkeypatch, fips_flag=int(openssl_fips), md5_refused=openssl_fips) + monkeypatch.setattr(rust_openssl, "is_fips_enabled", lambda: cryptography_fips) + + with pytest.raises(FipsModeError, match=refusal) if refusal else nullcontext(): + enforce_fips_mode() + + +class _Database: + """Prisma's raw-query edge over a database whose only LiteLLM table besides LiteLLM_Config is the key table.""" + + def __init__(self, *, marked_fips: bool, holds_rows: bool, readable: bool = True) -> None: + self.marked_fips = marked_fips + self.holds_rows = holds_rows + self.readable = readable + + async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: + if not self.readable: + raise RuntimeError('relation "LiteLLM_Config" does not exist') + if "information_schema.tables" in query: + return [{"table_name": "LiteLLM_VerificationToken"}] + if '"LiteLLM_VerificationToken"' in query: + return [{"marked_fips": self.marked_fips, "holds_rows": self.holds_rows}] + return [{"?column?": 1}] if self.marked_fips else [] + + async def execute_raw(self, query: str, *args: object) -> int: + self.marked_fips = True + return 1 + + +class _Prisma: + def __init__(self, database: _Database) -> None: + self.writer_db: Final = database + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "fips_mode, database, refusal, marked_after", + [ + ("true", _Database(marked_fips=False, holds_rows=False), None, True), + ("true", _Database(marked_fips=True, holds_rows=True), None, True), + ("true", _Database(marked_fips=False, holds_rows=True), "already holds data written without it", False), + ("true", _Database(marked_fips=False, holds_rows=False, readable=False), "could not be checked", False), + ("false", _Database(marked_fips=True, holds_rows=True), "set up with LITELLM_FIPS_MODE=true", True), + ("false", _Database(marked_fips=False, holds_rows=True), None, False), + ("false", _Database(marked_fips=True, holds_rows=True, readable=False), None, True), + ], + ids=[ + "fips-on-new-database", + "fips-on-fips-database", + "fips-on-used-database", + "fips-on-unreadable-database", + "fips-off-fips-database", + "fips-off-used-database", + "fips-off-unreadable-database", + ], +) +async def test_enforce_fips_database( + monkeypatch: pytest.MonkeyPatch, fips_mode: str, database: _Database, refusal: str | None, marked_after: bool +) -> None: + monkeypatch.setenv("LITELLM_FIPS_MODE", fips_mode) + + with pytest.raises(FipsModeError, match=refusal) if refusal else nullcontext(): + await enforce_fips_database(FipsRepository(_Prisma(database)), database_is_configured=True) + + assert database.marked_fips is marked_after + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "fips_mode, database_is_configured, refusal", + [("true", True, "can't be reached"), ("true", False, None), ("false", True, None)], + ids=["fips-on-unreachable", "fips-on-no-database", "fips-off-unreachable"], +) +async def test_enforce_fips_database_without_a_connected_database( + monkeypatch: pytest.MonkeyPatch, fips_mode: str, database_is_configured: bool, refusal: str | None +) -> None: + monkeypatch.setenv("LITELLM_FIPS_MODE", fips_mode) + + with pytest.raises(FipsModeError, match=refusal) if refusal else nullcontext(): + await enforce_fips_database(None, database_is_configured=database_is_configured) diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 3647146eb4e..99e78ae6385 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -10,7 +10,7 @@ import subprocess import sys import time import types -from collections.abc import Callable +from collections.abc import Callable, Mapping, Sequence from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Final @@ -1790,6 +1790,38 @@ async def test_proxy_startup_refuses_an_unsafe_master_key_even_when_the_database assert ("could not be checked" in announced[0]) == key_can_have_encrypted_the_database +class _DatabaseSetUpInFipsMode: + async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: + return [{"?column?": 1}] + + async def execute_raw(self, query: str, *args: object) -> int: + raise AssertionError("a refused boot must not write to the database") + + +class _PrismaClientOnAFipsDatabase: + def __init__(self) -> None: + self.writer_db: Final = _DatabaseSetUpInFipsMode() + + async def aclose(self) -> None: + pass + + +@pytest.mark.asyncio +async def test_proxy_startup_refuses_a_database_set_up_in_fips_mode_when_fips_mode_is_off( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + from litellm.proxy.common_utils.fips import FipsModeError + from litellm.proxy.proxy_server import proxy_startup_event + + _boot_with_general_settings(monkeypatch, tmp_path, {"master_key": "sk-a-safe-master-key"}) + monkeypatch.delenv("LITELLM_FIPS_MODE", raising=False) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _PrismaClientOnAFipsDatabase()) + + with pytest.raises(FipsModeError, match="set up with LITELLM_FIPS_MODE=true"): + async with proxy_startup_event(FastAPI()): + pass + + class _DatabaseWithOneStoredCredential: def __init__(self, ciphertext): self._ciphertext = ciphertext