mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
Merge 2277895897 into a3a52466a5
This commit is contained in:
devin-ai-integration[bot]
GitHub
commit
d7a09e543b
7 changed files with 1094 additions and 39 deletions
|
|
@ -359,28 +359,18 @@ MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY: Final = "litellm.model_group_alias_resolve
|
|||
REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged"
|
||||
REALTIME_SESSION_FAILURE_LOGGED_KEY: Final = "realtime_session_failure_logged"
|
||||
|
||||
# SSL/TLS cipher configuration for faster handshakes
|
||||
# Strategy: Strongly prefer fast modern ciphers, but allow fallback to commonly supported ones
|
||||
# This balances performance with broad compatibility
|
||||
FIPS_SSL_CIPHERS: Final = os.getenv(
|
||||
"LITELLM_SSL_CIPHERS",
|
||||
"TLS_AES_256_GCM_SHA384:TLS_AES_128_GCM_SHA256:"
|
||||
"ECDHE-RSA-AES256-GCM-SHA384:ECDHE-RSA-AES128-GCM-SHA256:"
|
||||
"ECDHE-ECDSA-AES256-GCM-SHA384:ECDHE-ECDSA-AES128-GCM-SHA256",
|
||||
)
|
||||
DEFAULT_SSL_CIPHERS: Final = os.getenv(
|
||||
"LITELLM_SSL_CIPHERS",
|
||||
# Priority 1: TLS 1.3 ciphers (fastest, ~50ms handshake)
|
||||
"TLS_AES_256_GCM_SHA384:" # Fastest observed in testing
|
||||
"TLS_AES_128_GCM_SHA256:" # Slightly faster than 256-bit
|
||||
"TLS_CHACHA20_POLY1305_SHA256:" # Fast on ARM/mobile
|
||||
# Priority 2: TLS 1.2 ECDHE+GCM (fast, ~100ms handshake, widely supported)
|
||||
"ECDHE-RSA-AES256-GCM-SHA384:"
|
||||
"ECDHE-RSA-AES128-GCM-SHA256:"
|
||||
"ECDHE-ECDSA-AES256-GCM-SHA384:"
|
||||
"ECDHE-ECDSA-AES128-GCM-SHA256:"
|
||||
# Priority 3: Additional modern ciphers (good balance)
|
||||
"ECDHE-RSA-CHACHA20-POLY1305:"
|
||||
"ECDHE-ECDSA-CHACHA20-POLY1305:"
|
||||
# Priority 4: Widely compatible fallbacks (slower but universally supported)
|
||||
"ECDHE-RSA-AES256-SHA384:" # Common fallback
|
||||
"ECDHE-RSA-AES128-SHA256:" # Very widely supported
|
||||
"AES256-GCM-SHA384:" # Non-PFS fallback (compatibility)
|
||||
"AES128-GCM-SHA256", # Last resort (maximum compatibility)
|
||||
"TLS_AES_256_GCM_SHA384:TLS_AES_128_GCM_SHA256:"
|
||||
"ECDHE-RSA-AES256-GCM-SHA384:ECDHE-RSA-AES128-GCM-SHA256:"
|
||||
"ECDHE-ECDSA-AES256-GCM-SHA384:ECDHE-ECDSA-AES128-GCM-SHA256:"
|
||||
"ECDHE-RSA-AES256-SHA384:ECDHE-RSA-AES128-SHA256",
|
||||
)
|
||||
|
||||
########### v2 Architecture constants for managing writing updates to the database ###########
|
||||
|
|
|
|||
|
|
@ -36,6 +36,7 @@ from litellm.constants import (
|
|||
AIOHTTP_TTL_DNS_CACHE,
|
||||
COMPLETION_HTTP_FALLBACK_SECONDS,
|
||||
DEFAULT_SSL_CIPHERS,
|
||||
FIPS_SSL_CIPHERS,
|
||||
HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS,
|
||||
)
|
||||
from litellm.litellm_core_utils.logging_utils import track_llm_api_timing
|
||||
|
|
@ -274,15 +275,22 @@ def _prepare_request_data_and_content(
|
|||
|
||||
|
||||
# Cache for SSL contexts to avoid creating duplicate contexts with the same configuration
|
||||
# Key: tuple of (cafile, ssl_security_level, ssl_ecdh_curve)
|
||||
# Key: tuple of (cafile, ssl_security_level, ssl_ecdh_curve, ssl_ciphers)
|
||||
# Value: ssl.SSLContext
|
||||
_ssl_context_cache: Final[dict[tuple[str | None, str | None, str | None], ssl.SSLContext]] = {}
|
||||
_ssl_context_cache: Final[dict[tuple[str | None, str | None, str | None, str], ssl.SSLContext]] = {}
|
||||
|
||||
|
||||
def _outbound_ssl_ciphers() -> str:
|
||||
from litellm.proxy.common_utils.fips import is_fips_mode
|
||||
|
||||
return FIPS_SSL_CIPHERS if is_fips_mode() else DEFAULT_SSL_CIPHERS
|
||||
|
||||
|
||||
def _create_ssl_context(
|
||||
cafile: str | None,
|
||||
ssl_security_level: str | None,
|
||||
ssl_ecdh_curve: str | None,
|
||||
ssl_ciphers: str,
|
||||
) -> ssl.SSLContext:
|
||||
"""
|
||||
Create an SSL context with the given configuration.
|
||||
|
|
@ -299,9 +307,7 @@ def _create_ssl_context(
|
|||
# User provided custom cipher configuration (e.g., via SSL_SECURITY_LEVEL env var)
|
||||
custom_ssl_context.set_ciphers(ssl_security_level)
|
||||
else:
|
||||
# Use optimized cipher list that strongly prefers fast ciphers
|
||||
# but falls back to widely compatible ones
|
||||
custom_ssl_context.set_ciphers(DEFAULT_SSL_CIPHERS)
|
||||
custom_ssl_context.set_ciphers(ssl_ciphers)
|
||||
|
||||
# Configure ECDH curve for key exchange (e.g., to disable PQC and improve performance)
|
||||
# Set SSL_ECDH_CURVE env var or litellm.ssl_ecdh_curve to 'X25519' to disable PQC
|
||||
|
|
@ -416,15 +422,15 @@ def get_ssl_configuration(
|
|||
cafile = certifi.where()
|
||||
|
||||
if ssl_verify is not False:
|
||||
# Create cache key from configuration parameters
|
||||
cache_key: Final = (cafile, ssl_security_level, ssl_ecdh_curve)
|
||||
ssl_ciphers: Final = _outbound_ssl_ciphers()
|
||||
cache_key: Final = (cafile, ssl_security_level, ssl_ecdh_curve, ssl_ciphers)
|
||||
|
||||
# Check if we have a cached SSL context for this configuration
|
||||
if cache_key not in _ssl_context_cache:
|
||||
_ssl_context_cache[cache_key] = _create_ssl_context(
|
||||
cafile=cafile,
|
||||
ssl_security_level=ssl_security_level,
|
||||
ssl_ecdh_curve=ssl_ecdh_curve,
|
||||
ssl_ciphers=ssl_ciphers,
|
||||
)
|
||||
|
||||
# Return the cached SSL context
|
||||
|
|
|
|||
|
|
@ -1,13 +1,23 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import datetime
|
||||
import ipaddress
|
||||
import json
|
||||
import ssl
|
||||
import threading
|
||||
import time
|
||||
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
|
||||
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
|
||||
|
||||
|
||||
def write_self_signed_cert(cert_dir: Path, names: tuple[str, ...] = ("localhost",)) -> tuple[Path, Path]:
|
||||
|
|
@ -46,3 +56,186 @@ def server_context(cert_file: Path, key_file: Path) -> ssl.SSLContext:
|
|||
context: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
|
||||
context.load_cert_chain(certfile=cert_file, keyfile=key_file)
|
||||
return context
|
||||
|
||||
|
||||
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
|
||||
|
||||
@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 _response_body(identifier: str, model: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": identifier,
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": "completed",
|
||||
"model": model,
|
||||
"parallel_tool_calls": True,
|
||||
"tool_choice": "auto",
|
||||
"tools": [],
|
||||
"top_p": 1.0,
|
||||
"temperature": 1.0,
|
||||
"truncation": "disabled",
|
||||
"store": True,
|
||||
"metadata": {},
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
"id": "msg_x",
|
||||
"role": "assistant",
|
||||
"status": "completed",
|
||||
"content": [{"type": "output_text", "text": "tls ok", "annotations": []}],
|
||||
}
|
||||
],
|
||||
"usage": {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7},
|
||||
}
|
||||
|
||||
|
||||
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)))
|
||||
if user is not None and "slow" in user:
|
||||
time.sleep(5)
|
||||
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 self.path.endswith("/responses"):
|
||||
response_id: Final = f"resp_{user or ''}"
|
||||
responses_payload: Final = json.dumps(_response_body(response_id, model)).encode()
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(responses_payload)))
|
||||
self.end_headers()
|
||||
self.wfile.write(responses_payload)
|
||||
return
|
||||
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: tuple[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 = (*self._requests, request)
|
||||
|
||||
def received(self) -> tuple[TlsPeerRequest, ...]:
|
||||
with self._lock:
|
||||
return self._requests
|
||||
|
||||
def start(self) -> None:
|
||||
server: Final = _CipherPeer(("127.0.0.1", self.port), self)
|
||||
context: Final = server_context(cert_file=self.cert_file, key_file=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()
|
||||
|
|
|
|||
463
tests/integration/configuration/test_tls_cipher_proxy.py
Normal file
463
tests/integration/configuration/test_tls_cipher_proxy.py
Normal file
|
|
@ -0,0 +1,463 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import psutil
|
||||
import pytest
|
||||
import yaml
|
||||
from anthropic import Anthropic
|
||||
from anthropic import APIStatusError as AnthropicAPIStatusError
|
||||
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 OwnedProxy, owned_proxy, owned_proxy_process
|
||||
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())
|
||||
)
|
||||
base_settings: Final = JSON_OBJECT.validate_python(base_config.get("litellm_settings", {}))
|
||||
config: Final = {
|
||||
**base_config,
|
||||
"litellm_settings": {**base_settings, "cache": False},
|
||||
"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]
|
||||
before: Final = len(tls_peers.gcm.received())
|
||||
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()[before:]
|
||||
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}"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def cipher_owned_proxy(
|
||||
tls_cert: tuple[Path, Path], tls_peers: Peers, tmp_path_factory: pytest.TempPathFactory
|
||||
) -> Iterator[OwnedProxy]:
|
||||
cert_file, _ = tls_cert
|
||||
directory: Final = tmp_path_factory.mktemp("tlsproxyworkers")
|
||||
base_config: Final = JSON_OBJECT.validate_python(
|
||||
yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
)
|
||||
base_settings: Final = JSON_OBJECT.validate_python(base_config.get("litellm_settings", {}))
|
||||
config: Final = {
|
||||
**base_config,
|
||||
"litellm_settings": {**base_settings, "cache": False},
|
||||
"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",
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
config_path: Final = directory / "tls_cipher_workers_config.yaml"
|
||||
config_path.write_text(yaml.safe_dump(config))
|
||||
with gateway_from_environment() as gateway:
|
||||
with owned_proxy_process(
|
||||
gateway, directory, {"SSL_CERT_FILE": str(cert_file)}, config=config_path, workers=2
|
||||
) as owned:
|
||||
yield owned
|
||||
|
||||
|
||||
def _responses_request(proxy: Gateway, model: str, user: str) -> str:
|
||||
response: Final = proxy.request("POST", "/v1/responses", {"model": model, "input": "tls", "user": user})
|
||||
assert response.status_code == 200, response.text
|
||||
return str(JSON_OBJECT.validate_json(response.text)["id"])
|
||||
|
||||
|
||||
def _messages_request(proxy: Gateway, model: str) -> str:
|
||||
with Anthropic(api_key=proxy.key, base_url=str(proxy.client.base_url), max_retries=0, timeout=60) as client:
|
||||
message: Final = client.messages.create(
|
||||
model=model, max_tokens=8, messages=[{"role": "user", "content": "tls"}]
|
||||
)
|
||||
return message.id
|
||||
|
||||
|
||||
def test_proxy_round_trips_responses_over_pfs_gcm_peer_and_writes_spend(
|
||||
cipher_proxy: Gateway, tls_peers: Peers
|
||||
) -> None:
|
||||
marker: Final = uuid.uuid4().hex[:8]
|
||||
user: Final = f"user19-resp-{marker}"
|
||||
before: Final = len(tls_peers.gcm.received())
|
||||
response_id: Final = _responses_request(cipher_proxy, "tls-peer", user)
|
||||
assert response_id.startswith("resp_"), response_id
|
||||
recorded: Final = tuple(
|
||||
record for record in tls_peers.gcm.received()[before:] if record.path.endswith("/responses")
|
||||
)
|
||||
assert [record.user for record in recorded] == [user], recorded
|
||||
assert all(record.cipher == PFS_GCM for record in recorded), recorded
|
||||
spend_rows: Final = eventually(
|
||||
lambda: read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (response_id,)),
|
||||
lambda values: len(values) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
assert spend_rows[0]["request_id"] == response_id
|
||||
|
||||
|
||||
def test_proxy_reports_ssl_error_for_responses_on_chacha20_only_peer(cipher_proxy: Gateway, tls_peers: Peers) -> None:
|
||||
before: Final = len(tls_peers.chacha.received())
|
||||
response: Final = cipher_proxy.request(
|
||||
"POST",
|
||||
"/v1/responses",
|
||||
{"model": "tls-chacha-peer", "input": "tls", "user": f"user20-{uuid.uuid4().hex[:8]}"},
|
||||
)
|
||||
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.chacha.received()) == before
|
||||
readiness: Final = cipher_proxy.client.get("/health/readiness")
|
||||
assert readiness.status_code == 200, readiness.text
|
||||
|
||||
|
||||
def test_proxy_round_trips_messages_over_pfs_gcm_peer(cipher_proxy: Gateway, tls_peers: Peers) -> None:
|
||||
before: Final = len(tls_peers.gcm.received())
|
||||
message_id: Final = _messages_request(cipher_proxy, "tls-peer")
|
||||
assert message_id, "messages response carried no id"
|
||||
recorded: Final = tls_peers.gcm.received()
|
||||
assert len(recorded) - before == 1, recorded[before:]
|
||||
assert recorded[-1].cipher == PFS_GCM, recorded[-1]
|
||||
|
||||
|
||||
def test_proxy_reports_ssl_error_for_messages_on_chacha20_only_peer(cipher_proxy: Gateway, tls_peers: Peers) -> None:
|
||||
before: Final = len(tls_peers.chacha.received())
|
||||
with pytest.raises(AnthropicAPIStatusError) as raised:
|
||||
_messages_request(cipher_proxy, "tls-chacha-peer")
|
||||
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
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _BurstOutcome:
|
||||
user: str | None
|
||||
status: int
|
||||
body: str
|
||||
finished_at: float
|
||||
|
||||
|
||||
def _chaos_request(proxy: Gateway, kind: str, index: int, marker: str) -> _BurstOutcome:
|
||||
user: Final = f"chaos-{kind}-{index}-{marker}"
|
||||
if kind in ("chat", "slow"):
|
||||
try:
|
||||
chat_response: Final = proxy.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": "tls-peer",
|
||||
"messages": [{"role": "user", "content": "tls"}],
|
||||
"user": user,
|
||||
"num_retries": 0,
|
||||
},
|
||||
)
|
||||
return _BurstOutcome(
|
||||
user=user, status=chat_response.status_code, body=chat_response.text, finished_at=time.monotonic()
|
||||
)
|
||||
except httpx.HTTPError:
|
||||
return _BurstOutcome(user=user, status=0, body="", finished_at=time.monotonic())
|
||||
if kind == "stream":
|
||||
try:
|
||||
stream_response: Final = proxy.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": "tls-peer",
|
||||
"messages": [{"role": "user", "content": "tls"}],
|
||||
"user": user,
|
||||
"stream": True,
|
||||
"num_retries": 0,
|
||||
},
|
||||
)
|
||||
return _BurstOutcome(
|
||||
user=user, status=stream_response.status_code, body=stream_response.text, finished_at=time.monotonic()
|
||||
)
|
||||
except httpx.HTTPError:
|
||||
return _BurstOutcome(user=user, status=0, body="", finished_at=time.monotonic())
|
||||
if kind == "responses":
|
||||
try:
|
||||
responses_response: Final = proxy.request(
|
||||
"POST",
|
||||
"/v1/responses",
|
||||
{"model": "tls-peer", "input": "tls", "user": user, "num_retries": 0},
|
||||
)
|
||||
return _BurstOutcome(
|
||||
user=user,
|
||||
status=responses_response.status_code,
|
||||
body=responses_response.text,
|
||||
finished_at=time.monotonic(),
|
||||
)
|
||||
except httpx.HTTPError:
|
||||
return _BurstOutcome(user=user, status=0, body="", finished_at=time.monotonic())
|
||||
try:
|
||||
messages_response: Final = proxy.request(
|
||||
"POST",
|
||||
"/v1/messages",
|
||||
{
|
||||
"model": "tls-peer",
|
||||
"max_tokens": 8,
|
||||
"messages": [{"role": "user", "content": f"tls {user}"}],
|
||||
"num_retries": 0,
|
||||
},
|
||||
headers={"anthropic-version": "2023-06-01"},
|
||||
)
|
||||
return _BurstOutcome(
|
||||
user=None,
|
||||
status=messages_response.status_code,
|
||||
body=messages_response.text,
|
||||
finished_at=time.monotonic(),
|
||||
)
|
||||
except httpx.HTTPError:
|
||||
return _BurstOutcome(user=None, status=0, body="", finished_at=time.monotonic())
|
||||
|
||||
|
||||
def test_proxy_burst_survives_peer_restart_with_every_outcome_recorded(cipher_proxy: Gateway, tls_peers: Peers) -> None:
|
||||
marker: Final = uuid.uuid4().hex[:8]
|
||||
kinds: Final = ("chat", "stream", "responses", "messages", "chat", "stream")
|
||||
before: Final = len(tls_peers.gcm.received())
|
||||
executor: Final = ThreadPoolExecutor(max_workers=30)
|
||||
with executor:
|
||||
futures: Final = tuple(
|
||||
executor.submit(_chaos_request, cipher_proxy, kind, index, marker) for index in range(5) for kind in kinds
|
||||
)
|
||||
eventually(
|
||||
lambda: sum(1 for future in futures if future.done() and future.result().status == 200),
|
||||
lambda successes: successes >= 3,
|
||||
seconds=60,
|
||||
)
|
||||
stopped_at: Final = time.monotonic()
|
||||
tls_peers.gcm.stop()
|
||||
tls_peers.gcm.start()
|
||||
finished: Final = tuple(future.result(timeout=120) for future in futures)
|
||||
assert len(finished) == 30, finished
|
||||
assert all(outcome.finished_at >= stopped_at for outcome in finished if outcome.status != 200), (
|
||||
stopped_at,
|
||||
finished,
|
||||
)
|
||||
assert not any(outcome.status == 0 for outcome in finished), finished
|
||||
records: Final = tls_peers.gcm.received()[before:]
|
||||
assert all(record.cipher == PFS_GCM for record in records), records
|
||||
submitted_users: Final = {outcome.user for outcome in finished if outcome.user is not None}
|
||||
successful_users: Final = {
|
||||
outcome.user for outcome in finished if outcome.status == 200 and outcome.user is not None
|
||||
}
|
||||
assert all(any(record.user == user for record in records) for user in successful_users), (successful_users, records)
|
||||
recorded_users: Final = {record.user for record in records if record.user is not None}
|
||||
assert recorded_users <= submitted_users, (recorded_users - submitted_users, records)
|
||||
assert all(sum(1 for record in records if record.user == user) <= 3 for user in submitted_users), (
|
||||
submitted_users,
|
||||
records,
|
||||
)
|
||||
anonymous: Final = sum(1 for record in records if record.user is None)
|
||||
messages_outcomes: Final = sum(1 for outcome in finished if outcome.user is None)
|
||||
messages_ok: Final = sum(1 for outcome in finished if outcome.user is None and outcome.status == 200)
|
||||
assert messages_ok <= anonymous <= messages_outcomes * 3, (anonymous, messages_ok, records)
|
||||
after_id: Final = _sync_completion(cipher_proxy, "tls-peer", f"chaos-after-{marker}")
|
||||
assert after_id == f"chatcmpl-chaos-after-{marker}"
|
||||
last_record: Final = tls_peers.gcm.received()[-1]
|
||||
assert last_record.cipher == PFS_GCM and last_record.user == f"chaos-after-{marker}", last_record
|
||||
readiness: Final = cipher_proxy.client.get("/health/readiness")
|
||||
assert readiness.status_code == 200, readiness.text
|
||||
|
||||
|
||||
def test_proxy_burst_survives_losing_one_worker(cipher_owned_proxy: OwnedProxy, tls_peers: Peers) -> None:
|
||||
proxy: Final = cipher_owned_proxy.gateway
|
||||
marker: Final = uuid.uuid4().hex[:8]
|
||||
executor: Final = ThreadPoolExecutor(max_workers=12)
|
||||
with executor:
|
||||
futures: Final = tuple(
|
||||
executor.submit(_chaos_request, proxy, "slow" if index % 2 else "chat", index, marker)
|
||||
for index in range(12)
|
||||
)
|
||||
eventually(
|
||||
lambda: sum(1 for future in futures if future.done() and future.result().status == 200),
|
||||
lambda successes: successes >= 2,
|
||||
seconds=60,
|
||||
)
|
||||
children: Final = psutil.Process(cipher_owned_proxy.process.pid).children()
|
||||
assert children, "two-worker proxy has no child processes to kill"
|
||||
killed_at: Final = time.monotonic()
|
||||
children[0].kill()
|
||||
finished: Final = tuple(future.result(timeout=120) for future in futures)
|
||||
assert len(finished) == 12, finished
|
||||
assert any(outcome.finished_at > killed_at for outcome in finished), finished
|
||||
post_kill_id: Final = _sync_completion(proxy, "tls-peer", f"post-kill-{marker}")
|
||||
assert post_kill_id == f"chatcmpl-post-kill-{marker}"
|
||||
|
||||
def _ready() -> bool:
|
||||
try:
|
||||
return proxy.client.get("/health/readiness").status_code == 200
|
||||
except httpx.TransportError:
|
||||
return False
|
||||
|
||||
eventually(_ready, lambda ready: ready, seconds=60)
|
||||
335
tests/integration/sdk/test_tls_cipher_wire.py
Normal file
335
tests/integration/sdk/test_tls_cipher_wire.py
Normal file
|
|
@ -0,0 +1,335 @@
|
|||
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=r"(?i)ssl|handshake"):
|
||||
_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=r"(?i)ssl|handshake"):
|
||||
_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=r"(?i)ssl|handshake"):
|
||||
_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=r"(?i)ssl|handshake"):
|
||||
_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=r"(?i)ssl|handshake"):
|
||||
_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=r"(?i)ssl|handshake"):
|
||||
_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=r"(?i)ssl|handshake"):
|
||||
_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=r"(?i)ssl|handshake"):
|
||||
_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]] = {}
|
||||
outcomes_lock: Final = threading.Lock()
|
||||
|
||||
def _record(index: int, outcome: BurstOutcome) -> None:
|
||||
with outcomes_lock:
|
||||
outcomes[index] = outcome
|
||||
|
||||
def _snapshot() -> tuple[BurstOutcome, ...]:
|
||||
with outcomes_lock:
|
||||
return tuple(outcomes.values())
|
||||
|
||||
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}")
|
||||
_record(index, BurstOutcome(f"user-{index}", time.monotonic(), response_id, None))
|
||||
except Exception as error:
|
||||
_record(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}"}
|
||||
)
|
||||
_record(
|
||||
index,
|
||||
BurstOutcome(
|
||||
f"user-{index}", time.monotonic(), str(JSON_OBJECT.validate_json(response.content)["id"]), None
|
||||
),
|
||||
)
|
||||
except Exception as error:
|
||||
_record(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: sum(1 for outcome in _snapshot() if outcome.error is None) >= 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)
|
||||
|
||||
finished: Final = _snapshot()
|
||||
assert len(finished) == BURST_SIZE
|
||||
_assert_burst_outcomes({int(outcome.user.removeprefix("user-")): outcome for outcome in finished})
|
||||
succeeded: Final = {outcome.user for outcome in finished if outcome.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(outcome.error is None and outcome.when < stopped_at for outcome in finished), (
|
||||
"no success before restart"
|
||||
)
|
||||
assert any(outcome.error is None and outcome.when > restarted_at for outcome in finished), (
|
||||
"no success after restart"
|
||||
)
|
||||
|
|
@ -322,6 +322,45 @@ def test_get_ssl_configuration_integration():
|
|||
assert ssl_context.verify_mode is not None
|
||||
|
||||
|
||||
def _tls12_cipher_names(cipher_string: str) -> set[str]:
|
||||
return {name for name in cipher_string.split(":") if name and not name.startswith("TLS_")}
|
||||
|
||||
|
||||
def _context_tls12_cipher_names(context: ssl.SSLContext) -> set[str]:
|
||||
return {cipher["name"] for cipher in context.get_ciphers() if cipher["protocol"] == "TLSv1.2"}
|
||||
|
||||
|
||||
def _supported_tls12_cipher_names() -> set[str]:
|
||||
unrestricted = ssl.create_default_context()
|
||||
return _context_tls12_cipher_names(unrestricted)
|
||||
|
||||
|
||||
def test_get_ssl_configuration_uses_fips_cipher_list_only_in_fips_mode(monkeypatch: pytest.MonkeyPatch):
|
||||
from litellm.constants import DEFAULT_SSL_CIPHERS, FIPS_SSL_CIPHERS
|
||||
from litellm.llms.custom_httpx import http_handler
|
||||
|
||||
monkeypatch.setattr(http_handler, "_ssl_context_cache", {})
|
||||
monkeypatch.delenv("LITELLM_FIPS_MODE", raising=False)
|
||||
monkeypatch.delenv("SSL_SECURITY_LEVEL", raising=False)
|
||||
monkeypatch.delenv("SSL_ECDH_CURVE", raising=False)
|
||||
|
||||
supported = _supported_tls12_cipher_names()
|
||||
|
||||
default_context = get_ssl_configuration()
|
||||
assert isinstance(default_context, ssl.SSLContext)
|
||||
assert _context_tls12_cipher_names(default_context) == (_tls12_cipher_names(DEFAULT_SSL_CIPHERS) & supported)
|
||||
assert get_ssl_configuration() is default_context
|
||||
|
||||
monkeypatch.setenv("LITELLM_FIPS_MODE", "true")
|
||||
fips_context = get_ssl_configuration()
|
||||
assert fips_context is not default_context
|
||||
assert isinstance(fips_context, ssl.SSLContext)
|
||||
fips_names = _context_tls12_cipher_names(fips_context)
|
||||
assert fips_names == (_tls12_cipher_names(FIPS_SSL_CIPHERS) & supported)
|
||||
assert all("-GCM-" in name for name in fips_names)
|
||||
assert get_ssl_configuration() is fips_context
|
||||
|
||||
|
||||
# Session Reuse Tests
|
||||
class MockClientSession:
|
||||
"""Mock ClientSession that is not callable"""
|
||||
|
|
|
|||
|
|
@ -1,15 +1,8 @@
|
|||
import ast
|
||||
import inspect
|
||||
import json
|
||||
from unittest import mock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
|
||||
import importlib
|
||||
import inspect
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm import constants
|
||||
|
|
@ -70,6 +63,42 @@ def _build_constant_env_var_map() -> dict[str, str]:
|
|||
return env_var_map
|
||||
|
||||
|
||||
def _cipher_suites(cipher_string: str) -> list[str]:
|
||||
return [suite for suite in cipher_string.split(":") if suite]
|
||||
|
||||
|
||||
def test_default_ssl_ciphers_excludes_chacha20():
|
||||
assert all("CHACHA" not in suite for suite in _cipher_suites(constants.DEFAULT_SSL_CIPHERS))
|
||||
|
||||
|
||||
def test_default_ssl_ciphers_tls12_suites_all_offer_forward_secrecy():
|
||||
tls12_suites = [suite for suite in _cipher_suites(constants.DEFAULT_SSL_CIPHERS) if not suite.startswith("TLS_")]
|
||||
assert tls12_suites
|
||||
assert all(suite.startswith("ECDHE-") for suite in tls12_suites)
|
||||
|
||||
|
||||
def test_fips_ssl_ciphers_is_a_subset_of_default():
|
||||
assert set(_cipher_suites(constants.FIPS_SSL_CIPHERS)) <= set(_cipher_suites(constants.DEFAULT_SSL_CIPHERS))
|
||||
|
||||
|
||||
def test_fips_ssl_ciphers_tls12_suites_are_all_gcm():
|
||||
tls12_suites = [suite for suite in _cipher_suites(constants.FIPS_SSL_CIPHERS) if not suite.startswith("TLS_")]
|
||||
assert tls12_suites
|
||||
assert all("-GCM-" in suite for suite in tls12_suites)
|
||||
|
||||
|
||||
def test_ssl_ciphers_env_override_applies_to_both_lists(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setenv("LITELLM_SSL_CIPHERS", "ECDHE-RSA-AES128-GCM-SHA256")
|
||||
importlib.reload(constants)
|
||||
try:
|
||||
assert constants.DEFAULT_SSL_CIPHERS == "ECDHE-RSA-AES128-GCM-SHA256"
|
||||
assert constants.FIPS_SSL_CIPHERS == "ECDHE-RSA-AES128-GCM-SHA256"
|
||||
finally:
|
||||
monkeypatch.delenv("LITELLM_SSL_CIPHERS")
|
||||
importlib.reload(constants)
|
||||
assert constants.DEFAULT_SSL_CIPHERS != "ECDHE-RSA-AES128-GCM-SHA256"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("cli_value", "litellm_cli_value", "expected"),
|
||||
[
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue