From 29945f1b66ee9ae8ba393c2d7f2ee20d3b625a17 Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 24 Sep 2026 08:06:13 +0000 Subject: [PATCH] test(integration): cover outbound TLS cipher list at the handler and proxy edges Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/_support/tls.py | 194 +++++++++++ .../configuration/test_tls_cipher_proxy.py | 188 +++++++++++ tests/integration/sdk/test_tls_cipher_wire.py | 314 ++++++++++++++++++ 3 files changed, 696 insertions(+) create mode 100644 tests/integration/_support/tls.py create mode 100644 tests/integration/configuration/test_tls_cipher_proxy.py create mode 100644 tests/integration/sdk/test_tls_cipher_wire.py diff --git a/tests/integration/_support/tls.py b/tests/integration/_support/tls.py new file mode 100644 index 00000000000..1abc8273ffe --- /dev/null +++ b/tests/integration/_support/tls.py @@ -0,0 +1,194 @@ +from __future__ import annotations + +import datetime +import ipaddress +import json +import ssl +import threading +from collections.abc import Generator +from contextlib import contextmanager +from dataclasses import dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Final, cast + +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 pydantic import JsonValue, TypeAdapter + +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +def write_self_signed_cert(cert_dir: Path) -> tuple[Path, Path]: + key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + now: Final = datetime.datetime.now(datetime.timezone.utc) + name: Final = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "localhost")]) + cert: Final = ( + x509.CertificateBuilder() + .subject_name(name) + .issuer_name(name) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(now - datetime.timedelta(days=1)) + .not_valid_after(now + datetime.timedelta(days=7)) + .add_extension( + x509.SubjectAlternativeName([x509.DNSName("localhost"), x509.IPAddress(ipaddress.ip_address("127.0.0.1"))]), + critical=False, + ) + .sign(key, hashes.SHA256()) + ) + cert_file: Final = cert_dir / "cert.pem" + key_file: Final = cert_dir / "key.pem" + cert_file.write_bytes(cert.public_bytes(serialization.Encoding.PEM)) + key_file.write_bytes( + key.private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.TraditionalOpenSSL, + serialization.NoEncryption(), + ) + ) + return cert_file, key_file + + +@dataclass(frozen=True, slots=True) +class TlsPeerRequest: + path: str + user: str | None + cipher: str + + +def _completion_body(identifier: str, model: str) -> dict[str, JsonValue]: + return { + "id": identifier, + "object": "chat.completion", + "model": model, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "tls ok"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + + +def _chunk_body( + identifier: str, model: str, content: str, role: str | None, finish_reason: str | None +) -> dict[str, JsonValue]: + delta: dict[str, JsonValue] = {"content": content} + if role is not None: + delta["role"] = role + return { + "id": identifier, + "object": "chat.completion.chunk", + "model": model, + "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], + } + + +class _CipherPeer(ThreadingHTTPServer): + daemon_threads = True + allow_reuse_address = True + + def __init__(self, address: tuple[str, int], owner: TlsPeer) -> None: + super().__init__(address, _CipherPeerHandler) + self.owner = owner + + +def _negotiated_cipher(handler: BaseHTTPRequestHandler) -> str: + info: Final = cast(ssl.SSLSocket, handler.connection).cipher() + return info[0] if info else "" + + +class _CipherPeerHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_POST(self) -> None: + peer: Final = cast(_CipherPeer, self.server).owner + body: Final = JSON_OBJECT.validate_json(self.rfile.read(int(self.headers["Content-Length"]))) + user_value: Final = body.get("user") + user: Final = user_value if isinstance(user_value, str) else None + peer.record(TlsPeerRequest(path=self.path, user=user, cipher=_negotiated_cipher(self))) + identifier: Final = f"chatcmpl-{user or ''}" + model_value: Final = body.get("model") + model: Final = model_value if isinstance(model_value, str) else "gpt-4o-mini" + if body.get("stream") is True: + chunks: Final = ( + f"data: {json.dumps(_chunk_body(identifier, model, 'tls', 'assistant', None))}\n\n" + + f"data: {json.dumps(_chunk_body(identifier, model, ' ok', None, 'stop'))}\n\n" + + "data: [DONE]\n\n" + ).encode() + self.send_response(200) + self.send_header("Content-Type", "text/event-stream") + self.send_header("Content-Length", str(len(chunks))) + self.end_headers() + self.wfile.write(chunks) + return + payload: Final = json.dumps(_completion_body(identifier, model)).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def log_message(self, format: str, *args: object) -> None: + return + + +class TlsPeer: + def __init__(self, cipher: str, cert_file: Path, key_file: Path) -> None: + self.cipher = cipher + self.cert_file = cert_file + self.key_file = key_file + self.port = 0 + self._server: _CipherPeer | None = None + self._thread: threading.Thread | None = None + self._lock = threading.Lock() + self._requests: list[TlsPeerRequest] = [] + + @property + def url(self) -> str: + return f"https://127.0.0.1:{self.port}" + + def record(self, request: TlsPeerRequest) -> None: + with self._lock: + self._requests.append(request) + + def received(self) -> tuple[TlsPeerRequest, ...]: + with self._lock: + return tuple(self._requests) + + def start(self) -> None: + server: Final = _CipherPeer(("127.0.0.1", self.port), self) + context: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + context.load_cert_chain(str(self.cert_file), str(self.key_file)) + context.maximum_version = ssl.TLSVersion.TLSv1_2 + context.set_ciphers(self.cipher) + server.socket = context.wrap_socket(server.socket, server_side=True) + self.port = server.server_address[1] + self._server = server + self._thread = threading.Thread(target=server.serve_forever, daemon=True) + self._thread.start() + + def stop(self) -> None: + if self._server is None: + return + self._server.shutdown() + self._server.server_close() + if self._thread is not None: + self._thread.join(timeout=10) + self._server = None + self._thread = None + + +@contextmanager +def tls_peer(cipher: str, cert_file: Path, key_file: Path) -> Generator[TlsPeer, None, None]: + peer: Final = TlsPeer(cipher, cert_file, key_file) + peer.start() + try: + yield peer + finally: + peer.stop() diff --git a/tests/integration/configuration/test_tls_cipher_proxy.py b/tests/integration/configuration/test_tls_cipher_proxy.py new file mode 100644 index 00000000000..e118d858a86 --- /dev/null +++ b/tests/integration/configuration/test_tls_cipher_proxy.py @@ -0,0 +1,188 @@ +from __future__ import annotations + +import asyncio +import uuid +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import pytest +import yaml +from openai import APIStatusError, AsyncOpenAI, OpenAI +from pydantic import JsonValue, TypeAdapter + +from tests.integration._support.client import Gateway, eventually, gateway_from_environment +from tests.integration._support.database import read_rows +from tests.integration._support.process import owned_proxy +from tests.integration._support.tls import TlsPeer, tls_peer, write_self_signed_cert + +CHACHA20: Final = "ECDHE-RSA-CHACHA20-POLY1305" +NONPFS_GCM: Final = "AES256-GCM-SHA384" +PFS_GCM: Final = "ECDHE-RSA-AES256-GCM-SHA384" +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +@dataclass(frozen=True, slots=True) +class Peers: + gcm: TlsPeer + chacha: TlsPeer + nonpfs: TlsPeer + + +@pytest.fixture(scope="module") +def tls_cert(tmp_path_factory: pytest.TempPathFactory) -> tuple[Path, Path]: + return write_self_signed_cert(tmp_path_factory.mktemp("tlscerts")) + + +@pytest.fixture(scope="module") +def tls_peers(tls_cert: tuple[Path, Path]) -> Iterator[Peers]: + cert_file, key_file = tls_cert + with ( + tls_peer(PFS_GCM, cert_file, key_file) as gcm, + tls_peer(CHACHA20, cert_file, key_file) as chacha, + tls_peer(NONPFS_GCM, cert_file, key_file) as nonpfs, + ): + yield Peers(gcm=gcm, chacha=chacha, nonpfs=nonpfs) + + +@pytest.fixture(scope="module") +def cipher_proxy( + tls_cert: tuple[Path, Path], tls_peers: Peers, tmp_path_factory: pytest.TempPathFactory +) -> Iterator[Gateway]: + cert_file, _ = tls_cert + directory: Final = tmp_path_factory.mktemp("tlsproxy") + base_config: Final = JSON_OBJECT.validate_python( + yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + ) + config: Final = { + **base_config, + "model_list": [ + { + "model_name": "tls-peer", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": f"{tls_peers.gcm.url}/v1", + "api_key": "sk-peer", + }, + }, + { + "model_name": "tls-chacha-peer", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": f"{tls_peers.chacha.url}/v1", + "api_key": "sk-peer", + }, + }, + { + "model_name": "tls-nonpfs-peer", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": f"{tls_peers.nonpfs.url}/v1", + "api_key": "sk-peer", + }, + }, + ], + } + config_path: Final = directory / "tls_cipher_config.yaml" + config_path.write_text(yaml.safe_dump(config)) + with gateway_from_environment() as gateway: + with owned_proxy(gateway, directory, {"SSL_CERT_FILE": str(cert_file)}, config=config_path) as candidate: + yield candidate + + +def _sync_completion(proxy: Gateway, model: str, user: str, *, stream: bool = False) -> str: + with OpenAI(api_key=proxy.key, base_url=f"{proxy.client.base_url}/v1", max_retries=0, timeout=60) as client: + if stream: + chunks: Final = list( + client.chat.completions.create( + model=model, messages=[{"role": "user", "content": "tls"}], user=user, stream=True + ) + ) + assert chunks, "stream produced no chunks" + assert all(chunk.id == f"chatcmpl-{user}" for chunk in chunks), chunks[0].id + return str(chunks[0].id) + response: Final = client.chat.completions.create( + model=model, messages=[{"role": "user", "content": "tls"}], user=user + ) + return str(response.id) + + +async def _async_completion(proxy: Gateway, model: str, user: str, *, stream: bool = False) -> str: + async with AsyncOpenAI( + api_key=proxy.key, base_url=f"{proxy.client.base_url}/v1", max_retries=0, timeout=60 + ) as client: + if stream: + chunk_ids: Final = [ + chunk.id + async for chunk in await client.chat.completions.create( + model=model, messages=[{"role": "user", "content": "tls"}], user=user, stream=True + ) + ] + assert chunk_ids, "stream produced no chunks" + assert all(chunk_id == f"chatcmpl-{user}" for chunk_id in chunk_ids), chunk_ids[0] + return str(chunk_ids[0]) + response: Final = await client.chat.completions.create( + model=model, messages=[{"role": "user", "content": "tls"}], user=user + ) + return str(response.id) + + +def test_proxy_round_trips_over_pfs_gcm_peer_and_writes_spend(cipher_proxy: Gateway, tls_peers: Peers) -> None: + marker: Final = uuid.uuid4().hex[:8] + sync_id: Final = _sync_completion(cipher_proxy, "tls-peer", f"user13-sync-{marker}") + stream_id: Final = _sync_completion(cipher_proxy, "tls-peer", f"user13-stream-{marker}", stream=True) + async_id: Final = asyncio.run(_async_completion(cipher_proxy, "tls-peer", f"user13-async-{marker}")) + async_stream_id: Final = asyncio.run( + _async_completion(cipher_proxy, "tls-peer", f"user13-astream-{marker}", stream=True) + ) + users: Final = {f"user13-{leg}-{marker}" for leg in ("sync", "stream", "async", "astream")} + recorded: Final = tls_peers.gcm.received() + assert {record.user for record in recorded} == users, recorded + assert all(record.cipher == PFS_GCM for record in recorded), recorded + spend_rows: Final = tuple( + eventually( + lambda response_id=response_id: read_rows( + 'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (response_id,) + ), + lambda values: len(values) == 1, + seconds=70, + ) + for response_id in (sync_id, stream_id, async_id, async_stream_id) + ) + assert [rows[0]["request_id"] for rows in spend_rows] == [sync_id, stream_id, async_id, async_stream_id] + + +def test_proxy_reports_ssl_error_for_chacha20_only_peer(cipher_proxy: Gateway, tls_peers: Peers) -> None: + marker: Final = uuid.uuid4().hex[:8] + before: Final = len(tls_peers.chacha.received()) + with pytest.raises(APIStatusError) as raised: + _sync_completion(cipher_proxy, "tls-chacha-peer", f"user14-{marker}") + assert raised.value.status_code >= 400 + assert "ssl" in str(raised.value).lower() or "connection error" in str(raised.value).lower(), str(raised.value) + assert len(tls_peers.chacha.received()) == before + readiness: Final = cipher_proxy.client.get("/health/readiness") + assert readiness.status_code == 200, readiness.text + recovered: Final = _sync_completion(cipher_proxy, "tls-peer", f"user14-ok-{marker}") + assert recovered == f"chatcmpl-user14-ok-{marker}" + + +def test_proxy_reports_ssl_error_for_non_pfs_peer(cipher_proxy: Gateway, tls_peers: Peers) -> None: + marker: Final = uuid.uuid4().hex[:8] + before: Final = len(tls_peers.nonpfs.received()) + response: Final = cipher_proxy.request( + "POST", + "/v1/chat/completions", + { + "model": "tls-nonpfs-peer", + "messages": [{"role": "user", "content": "tls"}], + "user": f"user15-{marker}", + }, + ) + assert response.status_code >= 400, response.text + assert "ssl" in response.text.lower() or "connection error" in response.text.lower(), response.text + assert len(tls_peers.nonpfs.received()) == before + readiness: Final = cipher_proxy.client.get("/health/readiness") + assert readiness.status_code == 200, readiness.text + recovered: Final = _sync_completion(cipher_proxy, "tls-peer", f"user15-ok-{marker}") + assert recovered == f"chatcmpl-user15-ok-{marker}" diff --git a/tests/integration/sdk/test_tls_cipher_wire.py b/tests/integration/sdk/test_tls_cipher_wire.py new file mode 100644 index 00000000000..62df651a293 --- /dev/null +++ b/tests/integration/sdk/test_tls_cipher_wire.py @@ -0,0 +1,314 @@ +from __future__ import annotations + +import asyncio +import os +import socket +import subprocess +import sys +import textwrap +import threading +import time +from collections.abc import Mapping +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Final + +import httpx +import pytest +from pydantic import JsonValue, TypeAdapter + +if TYPE_CHECKING: + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + +from tests.integration._support.client import eventually +from tests.integration._support.tls import TlsPeer, tls_peer, write_self_signed_cert + +CHACHA20: Final = "ECDHE-RSA-CHACHA20-POLY1305" +AES256_GCM: Final = "AES256-GCM-SHA384" +AES128_GCM: Final = "AES128-GCM-SHA256" +CBC_SHA256: Final = "ECDHE-RSA-AES128-SHA256" +CBC_SHA384: Final = "ECDHE-RSA-AES256-SHA384" +PFS_GCM: Final = "ECDHE-RSA-AES256-GCM-SHA384" +BURST_SIZE: Final = 30 +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + +OVERRIDE_EXCHANGE: Final = textwrap.dedent( + """ + import sys + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + handler = HTTPHandler(ssl_verify=sys.argv[2]) + try: + response = handler.post(f"{sys.argv[1]}/v1/chat/completions", json={"user": "u-override"}) + print(response.json()["id"]) + finally: + handler.close() + """ +) + + +@pytest.fixture(scope="module") +def tls_cert(tmp_path_factory: pytest.TempPathFactory) -> tuple[Path, Path]: + return write_self_signed_cert(tmp_path_factory.mktemp("tlscerts")) + + +@dataclass(frozen=True, slots=True) +class BurstOutcome: + user: str + when: float + response_id: str | None + error: Exception | None + + +def _sync_post(cert_file: Path, url: str, user: str) -> str: + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + handler: Final = HTTPHandler(ssl_verify=str(cert_file)) + try: + response: Final = handler.client.post(f"{url}/v1/chat/completions", json={"user": user}) + return str(JSON_OBJECT.validate_json(response.content)["id"]) + finally: + handler.close() + + +def _async_post(cert_file: Path, url: str, user: str) -> str: + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + async def _run() -> str: + handler: Final = AsyncHTTPHandler(ssl_verify=str(cert_file)) + try: + response: Final = await handler.client.post(f"{url}/v1/chat/completions", json={"user": user}) + return str(JSON_OBJECT.validate_json(response.content)["id"]) + finally: + await handler.close() + + return asyncio.run(_run()) + + +def _assert_served(peer: TlsPeer, user: str, response_id: str) -> None: + records: Final = peer.received() + assert response_id == f"chatcmpl-{user}" + assert [(record.user, record.cipher) for record in records] == [(user, peer.cipher)], records + + +def test_sync_handler_refuses_chacha20_only_peer(tls_cert: tuple[Path, Path]) -> None: + cert_file, key_file = tls_cert + with tls_peer(CHACHA20, cert_file, key_file) as peer: + with pytest.raises(httpx.ConnectError, match="SSL"): + _sync_post(cert_file, peer.url, "u-cell1") + assert peer.received() == (), peer.received() + + +def test_async_httpx_handler_refuses_chacha20_only_peer( + monkeypatch: pytest.MonkeyPatch, tls_cert: tuple[Path, Path] +) -> None: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + cert_file, key_file = tls_cert + with tls_peer(CHACHA20, cert_file, key_file) as peer: + with pytest.raises(httpx.ConnectError, match="SSL"): + _async_post(cert_file, peer.url, "u-cell2") + assert peer.received() == (), peer.received() + + +def test_async_aiohttp_handler_refuses_chacha20_only_peer(tls_cert: tuple[Path, Path]) -> None: + cert_file, key_file = tls_cert + with tls_peer(CHACHA20, cert_file, key_file) as peer: + with pytest.raises(httpx.ProtocolError, match="SSL"): + _async_post(cert_file, peer.url, "u-cell3") + assert peer.received() == (), peer.received() + + +def test_sync_handler_refuses_non_pfs_aes256_gcm_peer(tls_cert: tuple[Path, Path]) -> None: + cert_file, key_file = tls_cert + with tls_peer(AES256_GCM, cert_file, key_file) as peer: + with pytest.raises(httpx.ConnectError, match="SSL"): + _sync_post(cert_file, peer.url, "u-cell4") + assert peer.received() == (), peer.received() + + +def test_async_httpx_handler_refuses_non_pfs_aes128_gcm_peer( + monkeypatch: pytest.MonkeyPatch, tls_cert: tuple[Path, Path] +) -> None: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + cert_file, key_file = tls_cert + with tls_peer(AES128_GCM, cert_file, key_file) as peer: + with pytest.raises(httpx.ConnectError, match="SSL"): + _async_post(cert_file, peer.url, "u-cell5") + assert peer.received() == (), peer.received() + + +def test_sync_handler_accepts_cbc_peer_by_default(tls_cert: tuple[Path, Path]) -> None: + cert_file, key_file = tls_cert + with tls_peer(CBC_SHA256, cert_file, key_file) as peer: + _assert_served(peer, "u-cell6", _sync_post(cert_file, peer.url, "u-cell6")) + + +def test_sync_handler_refuses_cbc_peer_in_fips_mode( + monkeypatch: pytest.MonkeyPatch, tls_cert: tuple[Path, Path] +) -> None: + monkeypatch.setenv("LITELLM_FIPS_MODE", "true") + cert_file, key_file = tls_cert + with tls_peer(CBC_SHA256, cert_file, key_file) as peer: + with pytest.raises(httpx.ConnectError, match="SSL"): + _sync_post(cert_file, peer.url, "u-cell7") + assert peer.received() == (), peer.received() + + +def test_async_httpx_handler_refuses_cbc_sha384_peer_in_fips_mode( + monkeypatch: pytest.MonkeyPatch, tls_cert: tuple[Path, Path] +) -> None: + monkeypatch.setenv("LITELLM_FIPS_MODE", "true") + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + cert_file, key_file = tls_cert + with tls_peer(CBC_SHA384, cert_file, key_file) as peer: + with pytest.raises(httpx.ConnectError, match="SSL"): + _async_post(cert_file, peer.url, "u-cell8") + assert peer.received() == (), peer.received() + + +def test_handlers_accept_pfs_gcm_peer_in_default_and_fips_modes( + monkeypatch: pytest.MonkeyPatch, tls_cert: tuple[Path, Path] +) -> None: + cert_file, key_file = tls_cert + with tls_peer(PFS_GCM, cert_file, key_file) as peer: + served: Final = _sync_post(cert_file, peer.url, "u-cell9-sync") + assert served == "chatcmpl-u-cell9-sync" + served_async: Final = _async_post(cert_file, peer.url, "u-cell9-async") + assert served_async == "chatcmpl-u-cell9-async" + monkeypatch.setenv("LITELLM_FIPS_MODE", "true") + served_fips: Final = _sync_post(cert_file, peer.url, "u-cell9-fips") + assert served_fips == "chatcmpl-u-cell9-fips" + served_async_fips: Final = _async_post(cert_file, peer.url, "u-cell9-async-fips") + assert served_async_fips == "chatcmpl-u-cell9-async-fips" + assert {record.cipher for record in peer.received()} == {PFS_GCM} + + +def _override_exchange(url: str, cert_file: Path, overrides: Mapping[str, str]) -> str: + completed: Final = subprocess.run( + [sys.executable, "-P", "-c", OVERRIDE_EXCHANGE, url, str(cert_file)], + env={**os.environ, **overrides}, + capture_output=True, + text=True, + timeout=60, + check=False, + ) + assert completed.returncode == 0, completed.stderr + return completed.stdout.strip() + + +def test_ssl_ciphers_env_override_keeps_chacha20(tls_cert: tuple[Path, Path]) -> None: + cert_file, key_file = tls_cert + with tls_peer(CHACHA20, cert_file, key_file) as peer: + assert _override_exchange(peer.url, cert_file, {"LITELLM_SSL_CIPHERS": CHACHA20}) == "chatcmpl-u-override" + assert [record.cipher for record in peer.received()] == [CHACHA20] + + +def test_ssl_ciphers_override_wins_over_fips_mode(tls_cert: tuple[Path, Path]) -> None: + cert_file, key_file = tls_cert + with tls_peer(CHACHA20, cert_file, key_file) as peer: + result: Final = _override_exchange( + peer.url, cert_file, {"LITELLM_SSL_CIPHERS": CHACHA20, "LITELLM_FIPS_MODE": "true"} + ) + assert result == "chatcmpl-u-override" + assert [record.cipher for record in peer.received()] == [CHACHA20] + + +def test_cipher_list_follows_fips_mode_toggle_in_one_process( + monkeypatch: pytest.MonkeyPatch, tls_cert: tuple[Path, Path] +) -> None: + cert_file, key_file = tls_cert + with tls_peer(CBC_SHA256, cert_file, key_file) as peer: + assert _sync_post(cert_file, peer.url, "u-cell12-before") == "chatcmpl-u-cell12-before" + monkeypatch.setenv("LITELLM_FIPS_MODE", "true") + with pytest.raises(httpx.ConnectError, match="SSL"): + _sync_post(cert_file, peer.url, "u-cell12-fips") + monkeypatch.delenv("LITELLM_FIPS_MODE") + assert _sync_post(cert_file, peer.url, "u-cell12-after") == "chatcmpl-u-cell12-after" + assert [record.user for record in peer.received()] == ["u-cell12-before", "u-cell12-after"] + + +def _assert_burst_outcomes(outcomes: Mapping[int, BurstOutcome]) -> None: + for index, outcome in outcomes.items(): + if outcome.error is None: + assert outcome.response_id == f"chatcmpl-user-{index}", outcome + continue + message: str = str(outcome.error).lower() + assert "handshake" not in message and "cipher" not in message, ( + f"user-{index} failed with a cipher class error: {outcome.error!r}" + ) + assert isinstance(outcome.error, httpx.HTTPError), f"user-{index}: {outcome.error!r}" + + +def _peer_ready(port: int) -> bool: + try: + with socket.create_connection(("127.0.0.1", port), timeout=0.5): + return True + except OSError: + return False + + +def test_burst_survives_peer_restart_without_cipher_errors(tls_cert: tuple[Path, Path]) -> None: + cert_file, key_file = tls_cert + with tls_peer(PFS_GCM, cert_file, key_file) as peer: + second_wave: Final = threading.Event() + outcomes: Final[dict[int, BurstOutcome]] = {} + + def _one(index: int) -> None: + if index >= 10: + second_wave.wait(timeout=30) + try: + response_id: Final = _sync_post(cert_file, peer.url, f"user-{index}") + outcomes[index] = BurstOutcome(f"user-{index}", time.monotonic(), response_id, None) + except Exception as error: + outcomes[index] = BurstOutcome(f"user-{index}", time.monotonic(), None, error) + + async def _async_burst() -> None: + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + handler: Final = AsyncHTTPHandler(ssl_verify=str(cert_file)) + try: + await asyncio.to_thread(second_wave.wait, 30) + await asyncio.gather(*(_async_one(handler, index) for index in range(15, BURST_SIZE))) + finally: + await handler.close() + + async def _async_one(handler: AsyncHTTPHandler, index: int) -> None: + try: + response: Final = await handler.client.post( + f"{peer.url}/v1/chat/completions", json={"user": f"user-{index}"} + ) + outcomes[index] = BurstOutcome( + f"user-{index}", time.monotonic(), str(JSON_OBJECT.validate_json(response.content)["id"]), None + ) + except Exception as error: + outcomes[index] = BurstOutcome(f"user-{index}", time.monotonic(), None, error) + + def _async_thread() -> None: + asyncio.run(_async_burst()) + + async_thread: Final = threading.Thread(target=_async_thread, daemon=True) + with ThreadPoolExecutor(max_workers=15) as pool: + async_thread.start() + futures: Final = tuple(pool.submit(_one, index) for index in range(15)) + eventually(lambda: len(peer.received()) >= 3, lambda ready: ready, seconds=15) + stopped_at: Final = time.monotonic() + peer.stop() + peer.start() + eventually(lambda: _peer_ready(peer.port), lambda ready: ready, seconds=15) + restarted_at: Final = time.monotonic() + second_wave.set() + for future in futures: + future.result(timeout=60) + async_thread.join(timeout=60) + + assert len(outcomes) == BURST_SIZE + _assert_burst_outcomes(outcomes) + succeeded: Final = {o.user for o in outcomes.values() if o.error is None} + recorded: Final = peer.received() + assert sorted(record.user or "" for record in recorded) == sorted(succeeded) + assert all(record.cipher == PFS_GCM for record in recorded) + for user in succeeded: + assert sum(1 for record in recorded if record.user == user) == 1, user + assert any(o.error is None and o.when < stopped_at for o in outcomes.values()), "no success before restart" + assert any(o.error is None and o.when > restarted_at for o in outcomes.values()), "no success after restart"