From 1ac9cb2a1dcf17645ee440d23693845ad6e36e2f Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 15 Oct 2026 14:47:23 -0700 Subject: [PATCH] fix(proxy-extras): accept unencoded DB passwords in the v2 migration lock and redact its errors (#46592) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_proxy_extras/_redaction.py | 50 +++++ .../litellm_proxy_extras/migration_lock.py | 33 ++- .../litellm_proxy_extras/utils.py | 61 +----- .../test_v1_migration_error_logging.py | 74 +++++++ .../test_litellm_proxy_extras_utils.py | 193 +++++++++++++++++- 5 files changed, 358 insertions(+), 53 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/_redaction.py diff --git a/litellm-proxy-extras/litellm_proxy_extras/_redaction.py b/litellm-proxy-extras/litellm_proxy_extras/_redaction.py new file mode 100644 index 00000000000..152ae29b43b --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/_redaction.py @@ -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"(? 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}." diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 4a311331fc0..be785e2fd78 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -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"(? 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( diff --git a/tests/integration/database/test_v1_migration_error_logging.py b/tests/integration/database/test_v1_migration_error_logging.py index 43879d371b3..1fe78fb68f3 100644 --- a/tests/integration/database/test_v1_migration_error_logging.py +++ b/tests/integration/database/test_v1_migration_error_logging.py @@ -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: diff --git a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py index 75ea6e2a698..02bdf3429ea 100644 --- a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py +++ b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py @@ -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"