diff --git a/litellm/proxy/auth/login_utils.py b/litellm/proxy/auth/login_utils.py index e93bc2a56f2..ef0732221df 100644 --- a/litellm/proxy/auth/login_utils.py +++ b/litellm/proxy/auth/login_utils.py @@ -113,10 +113,13 @@ async def _rehash_password_if_needed(user_id: str, password: str, stored: str) - from litellm.proxy.proxy_server import prisma_client if prisma_client is not None: - await UserRepository(prisma_client).table.update( - where={"user_id": user_id}, - data={"password": hash_password(password)}, - ) + try: + await UserRepository(prisma_client).table.update_many( + where={"user_id": user_id, "password": stored}, + data={"password": hash_password(password)}, + ) + except Exception as e: # noqa: BLE001 # a failed rehash must never surface into the login + verbose_proxy_logger.warning("Login-time password rehash could not update user %s: %s", user_id, e) def get_ui_credentials(master_key: str | None) -> tuple[str, str]: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 42dfbf7e421..4330cf90d6d 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -7048,6 +7048,7 @@ def hash_token(token: str): PBKDF2_ITERATIONS: Final = 600_000 +PBKDF2_MAX_ITERATIONS: Final = 10_000_000 PBKDF2_PREFIX: Final = "pbkdf2:sha256:" SCRYPT_PREFIX: Final = "scrypt:" @@ -7081,11 +7082,12 @@ def _verify_pbkdf2(password: str, stored: str) -> bool: scheme, digest, iterations, salt, derived = stored.split(":") if (scheme, digest) != ("pbkdf2", "sha256"): return False - expected: Final = hashlib.pbkdf2_hmac( - "sha256", password.encode(), base64.b64decode(salt, validate=True), int(iterations) - ) + count: Final = int(iterations) + if not 1 <= count <= PBKDF2_MAX_ITERATIONS: + return False + expected: Final = hashlib.pbkdf2_hmac("sha256", password.encode(), base64.b64decode(salt, validate=True), count) return secrets.compare_digest(base64.b64decode(derived, validate=True), expected) - except (ValueError, binascii.Error, TypeError): + except (ValueError, binascii.Error, TypeError, OverflowError): return False diff --git a/tests/test_litellm/proxy/auth/test_login_utils.py b/tests/test_litellm/proxy/auth/test_login_utils.py index 5ec37195c58..83e8ee25e95 100644 --- a/tests/test_litellm/proxy/auth/test_login_utils.py +++ b/tests/test_litellm/proxy/auth/test_login_utils.py @@ -2365,14 +2365,14 @@ class TestRehashPasswordIfNeeded: stored = hashlib.sha256(password.encode()).hexdigest() mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_usertable.update = AsyncMock() + mock_prisma_client.db.litellm_usertable.update_many = AsyncMock() with patch( # test-quality-ok: the rehash writes to the database; faked so no DB is needed "litellm.proxy.proxy_server.prisma_client", mock_prisma_client ): await _rehash_password_if_needed("u-1", password, stored) - update_kwargs = mock_prisma_client.db.litellm_usertable.update.await_args.kwargs - assert update_kwargs["where"] == {"user_id": "u-1"} + update_kwargs = mock_prisma_client.db.litellm_usertable.update_many.await_args.kwargs + assert update_kwargs["where"] == {"user_id": "u-1", "password": stored} written = update_kwargs["data"]["password"] assert written.startswith("pbkdf2:sha256:") assert verify_password(password, written) @@ -2383,10 +2383,27 @@ class TestRehashPasswordIfNeeded: from litellm.proxy.utils import hash_password mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_usertable.update = AsyncMock() + mock_prisma_client.db.litellm_usertable.update_many = AsyncMock() with patch( # test-quality-ok: the rehash writes to the database; faked so no DB is needed "litellm.proxy.proxy_server.prisma_client", mock_prisma_client ): await _rehash_password_if_needed("u-1", "rehash-me-1", hash_password("rehash-me-1")) - mock_prisma_client.db.litellm_usertable.update.assert_not_called() + mock_prisma_client.db.litellm_usertable.update_many.assert_not_called() + + @pytest.mark.asyncio + async def test_a_failed_rehash_logs_a_warning_and_never_raises(self, caplog): + import logging + + from litellm._logging import verbose_proxy_logger + from litellm.proxy.auth.login_utils import _rehash_password_if_needed + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.update_many = AsyncMock(side_effect=RuntimeError("db down")) + with patch( # test-quality-ok: the rehash writes to the database; faked so no DB is needed + "litellm.proxy.proxy_server.prisma_client", mock_prisma_client + ): + with caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name): + await _rehash_password_if_needed("u-1", "rehash-me-1", "scrypt:stored") + + assert "Login-time password rehash could not update user u-1" in caplog.text diff --git a/tests/test_litellm/proxy/auth/test_password_hashing.py b/tests/test_litellm/proxy/auth/test_password_hashing.py index a02818546de..b2ac73e6e7d 100644 --- a/tests/test_litellm/proxy/auth/test_password_hashing.py +++ b/tests/test_litellm/proxy/auth/test_password_hashing.py @@ -106,6 +106,15 @@ class TestVerifyPasswordFormats: def test_malformed_pbkdf2_rows_return_false(self, stored): assert verify_password("test", stored) is False + @pytest.mark.parametrize( + "iterations", + (10_000_001, 10**30, 0), + ids=("above_max", "overflows_c_long", "zero"), + ) + def test_out_of_range_iteration_counts_return_false(self, iterations): + stored = _pbkdf2_row("test", 600_000).replace(":600000:", f":{iterations}:") + assert verify_password("test", stored) is False + def test_scrypt_invalid_base64_rejected(self): assert verify_password("test", "scrypt:not-valid-base64!!!") is False