fix(proxy): make login rehash conditional and non fatal, bound pbkdf2 iterations

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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]:

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
 
 

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

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