mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
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>
This commit is contained in:
parent
dea88965fb
commit
29945f1b66
3 changed files with 696 additions and 0 deletions
194
tests/integration/_support/tls.py
Normal file
194
tests/integration/_support/tls.py
Normal file
|
|
@ -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()
|
||||
188
tests/integration/configuration/test_tls_cipher_proxy.py
Normal file
188
tests/integration/configuration/test_tls_cipher_proxy.py
Normal file
|
|
@ -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}"
|
||||
314
tests/integration/sdk/test_tls_cipher_wire.py
Normal file
314
tests/integration/sdk/test_tls_cipher_wire.py
Normal file
|
|
@ -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"
|
||||
Loading…
Add table
Reference in a new issue