scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-16 04:29:31 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 139

fix(proxy-extras): accept unencoded DB passwords in the v2 migration lock and redact its errors (#46592)

Browse source 
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yassin Kortam 2026-10-15 14:47:23 -07:00 • committed by GitHub
parent 62af39aaf8
commit 1ac9cb2a1d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 358 additions and 53 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

50
litellm-proxy-extras/litellm_proxy_extras/_redaction.py Normal file
View file

@ -0,0 +1,50 @@
import functools
import os
import re
from collections.abc import Callable
from typing import Final
from urllib.parse import unquote, urlsplit
_REDACTED: Final = "REDACTED"
_PASSWORD_QUERY_KEYS: Final = frozenset(("password", "sslpassword"))
@functools.cache
def _secret_shape_redactor() -> Callable[[str], str]:
try:
from litellm._logging import redact_secrets
except ImportError:
return lambda text: text
return redact_secrets
def _url_passwords(url: str) -> frozenset[str]:
try:
parts: Final = urlsplit(url)
except ValueError:
return frozenset()
query_pairs: Final = tuple(pair.partition("=") for pair in parts.query.split("&"))
raw_query_passwords: Final = tuple(
value for key, separator, value in query_pairs if separator and key.lower() in _PASSWORD_QUERY_KEYS
)
raw_passwords: Final = ((parts.password,) if parts.password else ()) + raw_query_passwords
return frozenset(password for password in raw_passwords + tuple(map(unquote, raw_passwords)) if password)
def _configured_database_passwords() -> frozenset[str]:
database_url: Final = os.getenv("DATABASE_URL")
direct_url: Final = os.getenv("DIRECT_URL")
database_passwords: Final = _url_passwords(database_url) if database_url else frozenset()
direct_passwords: Final = _url_passwords(direct_url) if direct_url else frozenset()
return database_passwords | direct_passwords
def _redact_credentials(text: str) -> str:
"""Mask configured database passwords before passing the text to LiteLLM redaction."""
passwords: Final = sorted(_configured_database_passwords(), key=len, reverse=True)
alternation: Final = "|".join(re.escape(password) for password in passwords)
password_pattern: Final = (
re.compile(rf"(?<![A-Za-z0-9_])(?:{alternation})(?![A-Za-z0-9_])", re.IGNORECASE) if passwords else None
)
result: Final = password_pattern.sub(_REDACTED, text) if password_pattern is not None else text
return _secret_shape_redactor()(result)

33
litellm-proxy-extras/litellm_proxy_extras/migration_lock.py
View file

@ -1,4 +1,5 @@
import random
import re
import time
from collections.abc import Generator, Mapping
from contextlib import contextmanager
@ -7,6 +8,7 @@ from typing import TYPE_CHECKING, Final
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
from litellm_proxy_extras._logging import logger
from litellm_proxy_extras._redaction import _redact_credentials
from litellm_proxy_extras.prisma_toolchain import (
MIGRATION_LOCK_TIMEOUT_ENV_VAR,
migration_lock_timeout,
@ -14,6 +16,7 @@ from litellm_proxy_extras.prisma_toolchain import (
MIGRATION_LOCK_KEY: Final = int.from_bytes(b"llm_mig2", "big")
PRISMA_ADVISORY_LOCK_KEY: Final = 72707369
_SERVER_ANSWERED_MARKERS: Final = ("FATAL:", "SSL", "certificate")
if TYPE_CHECKING:
import psycopg
@ -49,6 +52,28 @@ def _try_lock(connection: "psycopg.Connection[tuple[object, ...]]", key: int = M
return row is not None and row.acquired
def _coordination_failure(database_url: str, exc: "psycopg.Error", *, connecting: bool) -> str:
import psycopg
prefix: Final = "Lost or could not establish v2 migration coordination with the database: "
raw: Final = str(exc)
unquoted: Final = re.sub(r"\"[^\"]*\"|'[^']*'", "", raw)
detail: Final = _redact_credentials(raw).strip()
if (
not connecting
or not isinstance(exc, psycopg.OperationalError)
or any(marker in unquoted for marker in _SERVER_ANSWERED_MARKERS)
):
return prefix + detail
parsed_url: Final = urlsplit(database_url)
try:
port: Final = parsed_url.port
except ValueError:
return prefix + f"P1001: Can't reach database server. {detail}"
host: Final = parsed_url.hostname or "localhost"
return prefix + f"P1001: Can't reach database server at `{host}:{port or 5432}`. {detail}"
@dataclass(frozen=True, slots=True)
class MigrationCoordinator:
connection: "psycopg.Connection[tuple[object, ...]]"
@ -81,7 +106,11 @@ def migration_lock(database_url: str) -> Generator[MigrationCoordinator, None, N
wait_seconds: Final = migration_lock_timeout()
deadline: Final = time.monotonic() + wait_seconds
try:
with psycopg.connect(database_url, connect_timeout=10, autocommit=True) as connection:
connection: Final = psycopg.connect(database_url, connect_timeout=10, autocommit=True)
except psycopg.Error as exc:
raise RuntimeError(_coordination_failure(database_url, exc, connecting=True)) from None
try:
with connection:
coordinator: Final = MigrationCoordinator(connection)
logger.info("Waiting for the v2 migration coordinator lock (up to %ss)", wait_seconds)
while time.monotonic() < deadline:
@ -94,7 +123,7 @@ def migration_lock(database_url: str) -> Generator[MigrationCoordinator, None, N
return
time.sleep(min(random.uniform(0.5, 1.5), max(0.0, deadline - time.monotonic())))
except psycopg.Error as exc:
raise RuntimeError(f"Lost or could not establish v2 migration coordination with the database: {exc}") from exc
raise RuntimeError(_coordination_failure(database_url, exc, connecting=False)) from None
raise RuntimeError(
f"Timed out waiting for another v2 migration resolver after {wait_seconds}s. "
f"Check the running migration or increase {MIGRATION_LOCK_TIMEOUT_ENV_VAR}."

61
litellm-proxy-extras/litellm_proxy_extras/utils.py
View file

@ -1,4 +1,3 @@
import functools
import glob
import os
import random
@ -13,10 +12,10 @@ from contextlib import AbstractContextManager, nullcontext
from dataclasses import dataclass, replace
from pathlib import Path
from typing import TYPE_CHECKING, Final, Optional, Protocol, Union
from urllib.parse import unquote, urlsplit
from litellm_proxy_extras import prisma_toolchain
from litellm_proxy_extras._logging import logger
from litellm_proxy_extras._redaction import _redact_credentials
from litellm_proxy_extras.migration_lock import held_migration_lock
from litellm_proxy_extras.prisma_toolchain import (
PRISMA_COMMAND_TIMEOUT_ENV_VAR,
@ -212,51 +211,6 @@ def _max_migration_timestamp(names) -> int:
return max(_migration_timestamp(n) for n in names)
_REDACTED: Final = "REDACTED"
_PASSWORD_QUERY_KEYS: Final = frozenset(("password", "sslpassword"))
@functools.cache
def _secret_shape_redactor() -> Callable[[str], str]:
try:
from litellm._logging import redact_secrets
except ImportError:
return lambda text: text
return redact_secrets
def _url_passwords(url: str) -> frozenset[str]:
try:
parts: Final = urlsplit(url)
except ValueError:
return frozenset()
query_pairs: Final = tuple(pair.partition("=") for pair in parts.query.split("&"))
raw_query_passwords: Final = tuple(
value for key, separator, value in query_pairs if separator and key.lower() in _PASSWORD_QUERY_KEYS
)
raw_passwords: Final = ((parts.password,) if parts.password else ()) + raw_query_passwords
return frozenset(password for password in raw_passwords + tuple(map(unquote, raw_passwords)) if password)
def _configured_database_passwords() -> frozenset[str]:
database_url: Final = os.getenv("DATABASE_URL")
direct_url: Final = os.getenv("DIRECT_URL")
database_passwords: Final = _url_passwords(database_url) if database_url else frozenset()
direct_passwords: Final = _url_passwords(direct_url) if direct_url else frozenset()
return database_passwords | direct_passwords
def _redact_credentials(text: str) -> str:
"""Mask configured database passwords before passing the text to LiteLLM redaction."""
passwords: Final = sorted(_configured_database_passwords(), key=len, reverse=True)
alternation: Final = "|".join(re.escape(password) for password in passwords)
password_pattern: Final = (
re.compile(rf"(?<![A-Za-z0-9_])(?:{alternation})(?![A-Za-z0-9_])", re.IGNORECASE) if passwords else None
)
result: Final = password_pattern.sub(_REDACTED, text) if password_pattern is not None else text
return _secret_shape_redactor()(result)
def _redacted_command(command: object) -> Union[str, tuple[str, ...], list[str]]:
if isinstance(command, tuple):
return tuple(_redact_credentials(str(argument)) for argument in command)
@ -789,16 +743,23 @@ class ProxyExtrasDBManager:
(connection_limit, pool_timeout, schema, pgbouncer, sslaccept, ...) and
translate Prisma's TLS params back, since libpq reads ``sslcert`` as a
client certificate where Prisma reads it as the CA."""
from urllib.parse import parse_qsl, quote, urlencode, urlparse, urlunparse
from urllib.parse import parse_qsl, quote, unquote, urlencode, urlparse, urlunparse
parsed: Final = urlparse(url)
netloc: Final = parsed.netloc
userinfo, separator, hostport = netloc.rpartition("@")
user, password_separator, password = userinfo.partition(":")
canonical_userinfo: Final = quote(unquote(user), safe="") + (
password_separator + quote(unquote(password), safe="") if password_separator else ""
)
canonical_netloc: Final = f"{canonical_userinfo}{separator}{hostport}" if separator else netloc
if not parsed.query:
return url
return url if canonical_netloc == netloc else urlunparse(parsed._replace(netloc=canonical_netloc))
pairs: Final = tuple(parse_qsl(parsed.query))
kept: Final = tuple((k, v) for k, v in pairs if k in LIBPQ_URL_PARAMS)
sslaccept: Final = next((v for k, v in pairs if k == "sslaccept"), None)
libpq_pairs: Final = ProxyExtrasDBManager._libpq_tls_params(kept, sslaccept)
return urlunparse(parsed._replace(query=urlencode(libpq_pairs, quote_via=quote)))
return urlunparse(parsed._replace(netloc=canonical_netloc, query=urlencode(libpq_pairs, quote_via=quote)))
@staticmethod
def _libpq_tls_params(

74
tests/integration/database/test_v1_migration_error_logging.py
View file

@ -8,6 +8,7 @@ import socketserver
import subprocess
import sys
import threading
import traceback
import uuid
from collections.abc import Callable, Iterator, Mapping
from contextlib import contextmanager, suppress
@ -18,6 +19,8 @@ from urllib.parse import quote, urlsplit, urlunsplit
import psycopg
import pytest
from litellm_proxy_extras.migration_lock import migration_lock
from litellm_proxy_extras.utils import ProxyExtrasDBManager
from psycopg import sql
from psycopg.types.json import Jsonb
@ -534,6 +537,77 @@ def test_v2_resolver_unreachable_database_exits_2_and_names_p1001(
_assert_no_password_fragments(output)
def test_v2_migration_lock_unreachable_database_names_p1001_and_redacts_password(
monkeypatch: pytest.MonkeyPatch,
) -> None:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener:
listener.bind(("127.0.0.1", 0))
port: Final = int(listener.getsockname()[1])
database_url: Final = f"postgresql://u:{PASSWORD}@127.0.0.1:{port}/litellm"
monkeypatch.setenv("DATABASE_URL", database_url)
lock_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(database_url)
prefix: Final = "Lost or could not establish v2 migration coordination with the database: "
with pytest.raises(RuntimeError) as error:
with migration_lock(lock_url):
pytest.fail("the closed database port unexpectedly connected")
message: Final = str(error.value)
assert message.startswith(prefix)
assert "P1001" in message
assert f"127.0.0.1:{port}" in message
assert [fragment for fragment in ("wr ong'pw9", "wr ong", "ong'pw9") if fragment in message] == []
with pytest.raises(RuntimeError) as raw_error:
with migration_lock(database_url):
pytest.fail("the raw password URL unexpectedly connected")
raw_message: Final = str(raw_error.value)
assert [fragment for fragment in ("wr ong'pw9", "wr ong", "ong'pw9") if fragment in raw_message] == []
raw_traceback: Final = "".join(traceback.format_exception(raw_error.value))
assert [fragment for fragment in ("wr ong'pw9", "wr ong", "ong'pw9") if fragment in raw_traceback] == []
def test_v2_resolver_applies_migrations_with_an_unencoded_password(
tmp_path: Path, request: pytest.FixtureRequest
) -> None:
with owned_database(PASSWORD) as encoded_url:
parsed: Final = urlsplit(encoded_url)
unencoded_url: Final = _database_url(
encoded_url,
parsed.username or "",
PASSWORD,
parsed.path.lstrip("/"),
encode_password=False,
)
completed: Final = run_v1_migrations(unencoded_url, tmp_path, {}, request.node.name, resolver="v2")
output: Final = completed.output
assert completed.returncode == 0, _safe_output(output)
assert error_lines(output) == (), _safe_output(output)
assert _applied_migrations(encoded_url) == SHIPPED_MIGRATIONS, _safe_output(output)
_assert_no_password_fragments(output)
def test_v2_resolver_wrong_unencoded_password_exits_2_without_p1001_or_password_fragments(
tmp_path: Path, request: pytest.FixtureRequest
) -> None:
correct_password: Final = f"correct-{uuid.uuid4().hex}"
with owned_database(correct_password) as encoded_url:
parsed: Final = urlsplit(encoded_url)
wrong_url: Final = _database_url(
encoded_url,
parsed.username or "",
PASSWORD,
parsed.path.lstrip("/"),
encode_password=False,
)
completed: Final = run_v1_migrations(wrong_url, tmp_path, {}, request.node.name, resolver="v2")
output: Final = completed.output
assert completed.returncode == 2, _safe_output(output)
assert "P1001" not in output, _safe_output(output)
_assert_no_password_fragments(output)
def test_v2_resolver_clean_database_applies_exactly_the_shipped_migrations(
tmp_path: Path, request: pytest.FixtureRequest
) -> None:

193
tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py
View file

@ -36,6 +36,122 @@ _MIGRATIONS_DIR = os.path.abspath(
)
def test_coordination_failure_does_not_misclassify_fatal_password_as_unreachable(
monkeypatch: pytest.MonkeyPatch,
) -> None:
import psycopg
from litellm_proxy_extras.migration_lock import _coordination_failure
database_url: Final = "postgresql://u:FATAL@127.0.0.1:5432/db"
monkeypatch.setenv("DATABASE_URL", database_url)
error: Final = psycopg.OperationalError(
'connection failed: connection to server at "127.0.0.1", port 5432 failed: '
'FATAL: password authentication failed for user "u"'
)
result: Final = _coordination_failure(database_url, error, connecting=True)
assert result.startswith("Lost or could not establish v2 migration coordination with the database: ")
assert "P1001" not in result
@pytest.mark.parametrize(
("message", "detail"),
[
(
'connection failed: connection to server at "127.0.0.1", port 5432 failed: SSL error: '
"certificate verify failed",
"certificate verify failed",
),
(
'connection failed: connection to server at "127.0.0.1", port 5432 failed: server certificate '
'for "db.internal" does not match host name "127.0.0.1"',
"does not match host name",
),
(
'connection failed: connection to server at "127.0.0.1", port 5432 failed: server does not support '
"SSL, but SSL was required",
"SSL was required",
),
],
)
def test_coordination_failure_does_not_label_tls_failures_as_unreachable(
monkeypatch: pytest.MonkeyPatch,
message: str,
detail: str,
) -> None:
import psycopg
from litellm_proxy_extras.migration_lock import _coordination_failure
database_url: Final = "postgresql://u:pw@127.0.0.1:5432/db"
monkeypatch.setenv("DATABASE_URL", database_url)
error: Final = psycopg.OperationalError(message)
result: Final = _coordination_failure(database_url, error, connecting=True)
assert result.startswith("Lost or could not establish v2 migration coordination with the database: ")
assert "P1001" not in result
assert detail in result
@pytest.mark.parametrize(
("database_url", "message", "host"),
[
(
"postgresql://u:pw@certificate-db.internal:5432/db",
'connection failed: connection to server at "certificate-db.internal" (10.0.0.5), port 5432 failed: '
"Connection refused",
"certificate-db.internal",
),
(
"postgresql://u:pw@SSL-primary:5432/db",
'connection failed: connection to server at "SSL-primary" (10.0.0.6), port 5432 failed: '
"Connection refused",
"ssl-primary",
),
(
"postgresql://u:pw@certificate-db.invalid:5432/db",
"connection failed: failed to resolve host 'certificate-db.invalid': [Errno -2] "
"Name or service not known",
"certificate-db.invalid",
),
],
)
def test_coordination_failure_labels_unreachable_hosts_named_like_tls_markers(
monkeypatch: pytest.MonkeyPatch,
database_url: str,
message: str,
host: str,
) -> None:
import psycopg
from litellm_proxy_extras.migration_lock import _coordination_failure
monkeypatch.setenv("DATABASE_URL", database_url)
error: Final = psycopg.OperationalError(message)
result: Final = _coordination_failure(database_url, error, connecting=True)
assert f"P1001: Can't reach database server at `{host}:5432`" in result
def test_coordination_failure_labels_connection_refused_as_unreachable(
monkeypatch: pytest.MonkeyPatch,
) -> None:
import psycopg
from litellm_proxy_extras.migration_lock import _coordination_failure
database_url: Final = "postgresql://u:pw@127.0.0.1:5432/db"
monkeypatch.setenv("DATABASE_URL", database_url)
error: Final = psycopg.OperationalError(
'connection failed: connection to server at "127.0.0.1", port 5432 failed: Connection refused'
)
result: Final = _coordination_failure(database_url, error, connecting=True)
assert result.startswith("Lost or could not establish v2 migration coordination with the database: ")
assert "P1001: Can't reach database server at `127.0.0.1:5432`" in result
def test_custom_prisma_dir(monkeypatch):
import tempfile
@ -681,8 +797,11 @@ class TestLensRenamePendingCheck:
with pytest.raises(psycopg.Error) as libpq:
psycopg.connect(database_url, connect_timeout=10, autocommit=True)
def raise_libpq_error(conninfo: str, *, connect_timeout: int, autocommit: bool) -> NoReturn:
raise libpq.value
with pytest.raises(RuntimeError) as err:
ProxyExtrasDBManager.raise_if_lens_rename_pending()
ProxyExtrasDBManager.raise_if_lens_rename_pending(connect=raise_libpq_error)
assert password not in str(err.value)
assert str(err.value).endswith(str(libpq.value).strip().replace(password, "REDACTED"))
@ -1463,6 +1582,78 @@ class TestStripPrismaQueryParams:
return dict(parse_qsl(urlparse(url).query))
@pytest.mark.parametrize(
"url,expected_user,expected_password,expected_host,expected_port,expected_sslmode",
(
(
"postgresql://u:wr ong'pw9@writer:5433/litellm",
"u",
"wr ong'pw9",
"writer",
"5433",
None,
),
(
"postgresql://u:wr ong'pw9@writer:5433/litellm?schema=x&sslmode=require",
"u",
"wr ong'pw9",
"writer",
"5433",
"require",
),
(
"postgresql://u:p%40ss%20w@writer:5433/litellm",
"u",
"p@ss w",
"writer",
"5433",
None,
),
(
"postgresql://u:p%z@writer:5433/litellm",
"u",
"p%z",
"writer",
"5433",
None,
),
(
"postgresql://u:wr ong'pw9@[::1]:5433/litellm",
"u",
"wr ong'pw9",
"::1",
"5433",
None,
),
),
)
def test_userinfo_is_encoded_for_libpq(
self,
url: str,
expected_user: str,
expected_password: str,
expected_host: str,
expected_port: str,
expected_sslmode: str | None,
) -> None:
from psycopg.conninfo import conninfo_to_dict
result: Final = ProxyExtrasDBManager._strip_prisma_query_params(url)
conninfo: Final = conninfo_to_dict(result)
assert conninfo["user"] == expected_user
assert conninfo["password"] == expected_password
assert conninfo["host"] == expected_host
assert conninfo["port"] == expected_port
assert conninfo["dbname"] == "litellm"
assert conninfo.get("sslmode") == expected_sslmode
assert "schema" not in conninfo
def test_empty_password_userinfo_keeps_its_separator(self) -> None:
url: Final = "postgresql://u:@writer/litellm"
assert ProxyExtrasDBManager._strip_prisma_query_params(url) == url
def test_prisma_ca_sslcert_becomes_sslrootcert_with_verify_full(self):
url = "postgresql://u:p@writer:5432/db?schema=public&sslmode=require&sslcert=/tmp/pinned.pem&sslaccept=strict"