scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 87

fix(proxy): pin multi-root CA bundle to the server's root before handing it to Prisma (#40428)

Browse source 
Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-09-09 11:19:29 -07:00 • committed by GitHub
parent f89e9ac749
commit 096984bfc2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 282 additions and 24 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

107
litellm/proxy/db/db_url_settings.py
View file

@ -34,12 +34,20 @@ writer's connection params (pool size, timeouts, pgbouncer mode) for the
ones the reader URL does not pin itself.
"""
import _ssl
import hashlib
import os
import socket
import ssl
import struct
import sys
import tempfile
import urllib.parse
from collections.abc import Mapping
from collections.abc import Callable, Mapping, Sequence
from functools import partial
from pathlib import Path
from types import MappingProxyType
from typing import Annotated, Final, cast
from typing import Annotated, Final, Protocol, TypeAlias, cast
from pydantic import AliasChoices, BeforeValidator, Field
from pydantic_settings import BaseSettings, SettingsConfigDict
@ -126,21 +134,100 @@ def add_missing_query_params(url: str, params: Mapping[str, str | int | float])
LIBPQ_VERIFY_SSLMODES: Final[frozenset[str]] = frozenset({"verify-ca", "verify-full"})
PEM_CERT_HEADER: Final = b"-----BEGIN CERTIFICATE-----"
PG_SSL_REQUEST: Final = struct.pack("!ii", 8, 80877103)
TLS_PROBE_TIMEOUT_SECONDS: Final = 10.0
RootCertResolver: TypeAlias = Callable[[str, str, int], str] # mutable-ok: Callable parameter syntax
def translate_libpq_ssl_params(url: str) -> str:
class _VerifiedChainSource(Protocol):
def get_verified_chain(self) -> Sequence[_ssl.Certificate] | None: ...
def _verified_chain_der(tls: ssl.SSLSocket) -> tuple[bytes, ...]:
if sys.version_info >= (3, 13):
return tuple(tls.get_verified_chain())
legacy: Final = cast( # cast-ok: the stub omits _sslobj, the C object has get_verified_chain since 3.10
"_VerifiedChainSource | None",
tls._sslobj, # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # public API only from 3.13
)
chain: Final = () if legacy is None else legacy.get_verified_chain() or ()
return tuple(cert.public_bytes(_ssl.ENCODING_DER) for cert in chain)
def _server_trust_anchor(cafile: str, host: str, port: int) -> bytes | None:
try:
context: Final = ssl.create_default_context(cafile=cafile)
with socket.create_connection((host, port), timeout=TLS_PROBE_TIMEOUT_SECONDS) as raw:
raw.sendall(PG_SSL_REQUEST)
if raw.recv(1) != b"S":
return None
with context.wrap_socket(raw, server_hostname=host) as tls:
chain: Final = _verified_chain_der(tls)
except (OSError, ValueError):
return None
return chain[-1] if chain else None
def pin_bundle_root(cert_path: str, host: str, port: int) -> str:
"""Reduce a multi-root CA bundle to the one root that verifies ``host``.
Prisma's ``sslcert`` loads a single PEM certificate (native-tls
``Certificate::from_pem``), so pointing it at a bundle such as the AWS RDS
global bundle trusts only the first of its 108 regional roots and the
handshake fails with "unable to get local issuer certificate" for every
other region. A single-certificate file is returned as is. For a bundle,
one verifying handshake (chain and hostname, whole bundle as trust store)
identifies the trust anchor the server actually chains to, which is
written to a single-certificate file for Prisma. If the probe fails the
bundle path is returned unchanged, so Prisma fails closed exactly as
before rather than trusting anything the bundle would not.
"""
try:
if Path(cert_path).read_bytes().count(PEM_CERT_HEADER) < 2:
return cert_path
except OSError:
return cert_path
root: Final = _server_trust_anchor(cert_path, host, port)
if root is None:
return cert_path
pinned: Final = Path(tempfile.gettempdir()) / f"litellm-sslcert-{hashlib.sha256(root).hexdigest()[:16]}.pem"
return str(pinned) if _replace_file(pinned, ssl.DER_cert_to_PEM_cert(root)) else cert_path
def _replace_file(target: Path, content: str) -> bool:
"""Write ``content`` to a private temp file and rename it over ``target``, so
readers never see a partial file and a symlink planted at ``target`` is
replaced rather than followed."""
try:
fd, staged = tempfile.mkstemp(dir=target.parent, prefix=f"{target.name}.")
except OSError:
return False
try:
with os.fdopen(fd, "w") as handle:
handle.write(content)
os.replace(staged, target)
except OSError:
Path(staged).unlink(missing_ok=True)
return False
return True
def translate_libpq_ssl_params(url: str, resolve_root_cert: RootCertResolver = pin_bundle_root) -> str:
"""Rewrite libpq's certificate-verification params into Prisma's dialect.
Prisma's engine only knows ``sslmode=disable|prefer|require``, ``sslcert``
(the CA bundle) and ``sslaccept=strict``. It silently discards
(a single CA certificate) and ``sslaccept=strict``. It silently discards
``sslrootcert`` and downgrades ``sslmode=verify-ca`` / ``verify-full`` to
``prefer``, so a URL copied from libpq / RDS docs connects over TLS with no
certificate check at all. ``verify-ca`` and ``verify-full`` both become
``require`` (Prisma has no CA-only mode), ``sslrootcert`` becomes
``sslcert``, and either one turns on ``sslaccept=strict`` (chain and
hostname), matching libpq where a root cert makes ``require`` verify.
Prisma params the operator pinned themselves win; anything else is left
untouched.
``sslcert`` (run through ``resolve_root_cert``, which pins a multi-root
bundle down to the server's root), and either one turns on
``sslaccept=strict`` (chain and hostname), matching libpq where a root
cert makes ``require`` verify. Prisma params the operator pinned
themselves win; anything else is left untouched.
"""
parsed: Final = urllib.parse.urlsplit(url)
pairs: Final = tuple(urllib.parse.parse_qsl(parsed.query, keep_blank_values=True))
@ -154,7 +241,9 @@ def translate_libpq_ssl_params(url: str) -> str:
if key != "sslrootcert"
)
root_cert: Final = tuple(
("sslcert", value) for key, value in pairs if key == "sslrootcert" and "sslcert" not in keys
("sslcert", resolve_root_cert(value, parsed.hostname or "", parsed.port or int(DEFAULT_POSTGRES_PORT)))
for key, value in pairs
if key == "sslrootcert" and "sslcert" not in keys
)
strict: Final = () if "sslaccept" in keys else (("sslaccept", "strict"),)
query: Final = urllib.parse.urlencode(translated + root_cert + strict)

199
tests/test_litellm/proxy/db/test_db_url_settings.py
View file

@ -11,15 +11,30 @@ clobber a pre-existing ``DATABASE_URL_READ_REPLICA``. A pre-existing
``DATABASE_URL`` (password auth) is likewise left untouched.
"""
import datetime
import hashlib
import os
import socket
import ssl
import tempfile
import threading
import urllib.parse
from collections.abc import Iterator
from dataclasses import dataclass
from pathlib import Path
from typing import Final
from unittest.mock import patch
import pytest
from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import ec
from pydantic import ValidationError
from litellm.proxy.db.db_url_settings import (
PG_SSL_REQUEST,
DatabaseURLSettings,
translate_libpq_ssl_params,
unsupported_db_scheme,
unsupported_db_scheme_message,
)
@ -381,9 +396,7 @@ def test_writer_password_is_percent_encoded(monkeypatch):
def test_writer_url_not_clobbered_when_already_set(monkeypatch):
"""An operator-pinned DATABASE_URL (e.g. helm's $(VAR) assembly) always
wins over the discrete fields."""
monkeypatch.setenv(
"DATABASE_URL", "postgresql://pinned:url@db.example.com:5432/litellm_db"
)
monkeypatch.setenv("DATABASE_URL", "postgresql://pinned:url@db.example.com:5432/litellm_db")
monkeypatch.setenv("DATABASE_HOST", "writer.example.com")
monkeypatch.setenv("DATABASE_USER", "litellm")
monkeypatch.setenv("DATABASE_NAME", "litellm_db")
@ -515,9 +528,7 @@ def test_apply_to_env_rejects_pinned_sqlite_direct_url(monkeypatch):
def test_apply_to_env_rejects_pinned_non_postgres_reader(monkeypatch):
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@writer.example.com:5432/db")
monkeypatch.setenv(
"DATABASE_URL_READ_REPLICA", "mysql://u:p@reader.example.com:3306/db"
)
monkeypatch.setenv("DATABASE_URL_READ_REPLICA", "mysql://u:p@reader.example.com:3306/db")
with pytest.raises(RuntimeError, match=r"DATABASE_URL_READ_REPLICA.*mysql"):
_apply()
@ -542,15 +553,11 @@ def test_reader_inherits_writer_connection_params(monkeypatch):
"DATABASE_URL",
"postgresql://u:p@writer.example.com:5432/db?connection_limit=3&pool_timeout=20&pgbouncer=true",
)
monkeypatch.setenv(
"DATABASE_URL_READ_REPLICA", "postgresql://u:p@reader.example.com:5432/db"
)
monkeypatch.setenv("DATABASE_URL_READ_REPLICA", "postgresql://u:p@reader.example.com:5432/db")
_apply()
query = urllib.parse.parse_qs(
urllib.parse.urlsplit(os.environ["DATABASE_URL_READ_REPLICA"]).query
)
query = urllib.parse.parse_qs(urllib.parse.urlsplit(os.environ["DATABASE_URL_READ_REPLICA"]).query)
assert query["connection_limit"] == ["3"]
assert query["pool_timeout"] == ["20"]
assert query["pgbouncer"] == ["true"]
@ -568,9 +575,7 @@ def test_reader_keeps_its_own_pinned_connection_params(monkeypatch):
_apply()
query = urllib.parse.parse_qs(
urllib.parse.urlsplit(os.environ["DATABASE_URL_READ_REPLICA"]).query
)
query = urllib.parse.parse_qs(urllib.parse.urlsplit(os.environ["DATABASE_URL_READ_REPLICA"]).query)
assert query["connection_limit"] == ["50"]
assert query["pool_timeout"] == ["20"]
@ -776,6 +781,170 @@ def test_libpq_verify_full_and_sslrootcert_become_prisma_strict_sslcert(monkeypa
}
def _issue_cert(
subject: str, issuer: x509.Certificate | None, issuer_key: ec.EllipticCurvePrivateKey | None, ca: bool
) -> tuple[x509.Certificate, ec.EllipticCurvePrivateKey]:
key: Final = ec.generate_private_key(ec.SECP256R1())
name: Final = x509.Name((x509.NameAttribute(x509.NameOID.COMMON_NAME, subject),))
now: Final = datetime.datetime.now(datetime.timezone.utc)
builder: Final = (
x509.CertificateBuilder()
.subject_name(name)
.issuer_name(issuer.subject if issuer else name)
.public_key(key.public_key())
.serial_number(x509.random_serial_number())
.not_valid_before(now - datetime.timedelta(minutes=5))
.not_valid_after(now + datetime.timedelta(days=1))
.add_extension(x509.BasicConstraints(ca=ca, path_length=None), critical=True)
.add_extension(x509.SubjectAlternativeName((x509.DNSName("localhost"),)), critical=False)
)
return builder.sign(issuer_key or key, hashes.SHA256()), key
def _pem(cert: x509.Certificate) -> bytes:
return cert.public_bytes(serialization.Encoding.PEM)
class _TlsPostgresStub:
"""Answers one libpq ``SSLRequest`` with ``S`` and serves ``leaf + intermediate``."""
def __init__(self, chain_pem: Path, key_pem: Path) -> None:
self.context: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
self.context.load_cert_chain(str(chain_pem), str(key_pem))
self.listener: Final = socket.create_server(("127.0.0.1", 0))
self.port: Final[int] = self.listener.getsockname()[1]
self.thread: Final = threading.Thread(target=self._serve, daemon=True)
self.thread.start()
def _serve(self) -> None:
with self.listener:
while True:
try:
conn: socket.socket = self.listener.accept()[0]
except OSError:
return
with conn:
try:
if conn.recv(8) == PG_SSL_REQUEST:
conn.sendall(b"S")
with self.context.wrap_socket(conn, server_side=True) as tls:
tls.recv(1)
except OSError:
continue
@dataclass(frozen=True, slots=True)
class _RdsLikePki:
bundle: Path
wrong_bundle: Path
root: Path
port: int
@pytest.fixture
def rds_like_pki(tmp_path: Path) -> Iterator[_RdsLikePki]:
"""An RDS-shaped trust setup: the server sends leaf + intermediate, the
bundle holds only self-signed roots, and the right root is not first."""
root, root_key = _issue_cert("Real Root CA", None, None, ca=True)
decoys: Final = tuple(_issue_cert(f"Decoy Root CA {i}", None, None, ca=True)[0] for i in range(3))
intermediate, intermediate_key = _issue_cert("Intermediate CA", root, root_key, ca=True)
leaf, leaf_key = _issue_cert("localhost", intermediate, intermediate_key, ca=False)
chain_pem: Final = tmp_path / "server-chain.pem"
chain_pem.write_bytes(_pem(leaf) + _pem(intermediate))
key_pem: Final = tmp_path / "server.key"
key_pem.write_bytes(
leaf_key.private_bytes(
serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption()
)
)
bundle: Final = tmp_path / "global-bundle.pem"
bundle.write_bytes(b"".join(_pem(decoy) for decoy in decoys) + _pem(root))
wrong_bundle: Final = tmp_path / "wrong-bundle.pem"
wrong_bundle.write_bytes(b"".join(_pem(decoy) for decoy in decoys))
root_pem: Final = tmp_path / "root.pem"
root_pem.write_bytes(_pem(root))
stub: Final = _TlsPostgresStub(chain_pem, key_pem)
yield _RdsLikePki(bundle=bundle, wrong_bundle=wrong_bundle, root=root_pem, port=stub.port)
stub.listener.close()
def _params(url: str) -> tuple[tuple[str, str], ...]:
return tuple(urllib.parse.parse_qsl(urllib.parse.urlsplit(url).query, keep_blank_values=True))
def test_multi_root_bundle_is_pinned_to_the_root_the_server_chains_to(
monkeypatch: pytest.MonkeyPatch, rds_like_pki: _RdsLikePki
):
"""Prisma's ``sslcert`` loads only the first certificate of the file, so
handing it the whole RDS bundle trusts one region's root and fails with
"unable to get local issuer certificate" everywhere else. The URL Prisma
receives must point at a single-certificate file holding the server's root."""
monkeypatch.setenv(
"DATABASE_URL",
f"postgresql://u:p@localhost:{rds_like_pki.port}/litellm_db?sslmode=verify-full&sslrootcert={rds_like_pki.bundle}",
)
_apply()
(sslmode, sslcert, sslaccept, _) = _params(os.environ["DATABASE_URL"])
assert (sslmode, sslaccept) == (("sslmode", "require"), ("sslaccept", "strict"))
assert sslcert[0] == "sslcert" and sslcert[1] != str(rds_like_pki.bundle)
assert Path(sslcert[1]).read_bytes() == rds_like_pki.root.read_bytes()
def test_pinned_root_replaces_a_planted_symlink_instead_of_writing_through_it(
monkeypatch: pytest.MonkeyPatch, rds_like_pki: _RdsLikePki, tmp_path: Path
):
"""The pinned file has a predictable name in a shared temp dir, so a symlink
planted there must not redirect the write onto its target."""
monkeypatch.setattr(tempfile, "tempdir", str(tmp_path))
root_der: Final = x509.load_pem_x509_certificate(rds_like_pki.root.read_bytes()).public_bytes(
serialization.Encoding.DER
)
pinned: Final = tmp_path / f"litellm-sslcert-{hashlib.sha256(root_der).hexdigest()[:16]}.pem"
victim: Final = tmp_path / "victim.txt"
victim.write_text("untouched")
pinned.symlink_to(victim)
monkeypatch.setenv(
"DATABASE_URL",
f"postgresql://u:p@localhost:{rds_like_pki.port}/litellm_db?sslmode=verify-full&sslrootcert={rds_like_pki.bundle}",
)
_apply()
assert ("sslcert", str(pinned)) in _params(os.environ["DATABASE_URL"])
assert victim.read_text() == "untouched"
assert not pinned.is_symlink() and pinned.read_bytes() == rds_like_pki.root.read_bytes()
def test_bundle_without_the_servers_root_is_passed_through_unchanged(
monkeypatch: pytest.MonkeyPatch, rds_like_pki: _RdsLikePki
):
"""Nothing in the bundle verifies the server, so no root is pinned and
Prisma keeps rejecting the connection instead of trusting a root the
operator never shipped."""
monkeypatch.setenv(
"DATABASE_URL",
f"postgresql://u:p@localhost:{rds_like_pki.port}/litellm_db"
f"?sslmode=verify-full&sslrootcert={rds_like_pki.wrong_bundle}",
)
_apply()
assert ("sslcert", str(rds_like_pki.wrong_bundle)) in _params(os.environ["DATABASE_URL"])
def test_root_cert_resolver_receives_the_urls_host_and_default_port():
def resolver(cert_path: str, host: str, port: int) -> str:
return f"/pinned/{host}/{port}{cert_path}"
url: Final = translate_libpq_ssl_params(
"postgresql://u:p@db.example.com/litellm_db?sslmode=verify-full&sslrootcert=/certs/bundle.pem", resolver
)
assert ("sslcert", "/pinned/db.example.com/5432/certs/bundle.pem") in _params(url)
def test_libpq_verify_ca_becomes_prisma_strict(monkeypatch):
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@db.example.com:5432/litellm_db?sslmode=verify-ca")