From 0a8423d77b7fe99e572d8ea5e923fcd05c6985a2 Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 17 Sep 2026 00:28:39 +0000 Subject: [PATCH] fix(proxy): keep the sign-in hold pool from refusing a correct password The held-attempt cap ran before the password check, so five parked wrong guesses from a blocked source turned the soft block into a lockout for the real user. The slot is now taken only after a wrong password, and the pool-full refusal carries the block's remaining time as Retry-After Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/login_throttle.py | 49 ++++++++++--------- .../proxy/auth/test_login_utils.py | 44 ++++++++++++++++- 2 files changed, 69 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/auth/login_throttle.py b/litellm/proxy/auth/login_throttle.py index d17e3fd0d53..a8899b9c675 100644 --- a/litellm/proxy/auth/login_throttle.py +++ b/litellm/proxy/auth/login_throttle.py @@ -281,25 +281,7 @@ class LoginThrottle: yield LoginAttempt(throttle=self, username=username, block=None) return slot: Final = keys.pair_block if block.scope == "user" else keys.source_block - held: Final = _HELD_ATTEMPTS.get(slot, 0) - if held >= MAX_HELD_ATTEMPTS_PER_KEY: - verbose_proxy_logger.warning( - "Admin UI sign-in refused: %s attempts already held for a blocked %s; username=%r source=%s", - held, - block.scope, - username, - self.client_ip, - ) - self.refuse(BLOCKED_ATTEMPT_HOLD_SECONDS) - _HELD_ATTEMPTS[slot] = held + 1 - try: - yield LoginAttempt(throttle=self, username=username, block=block) - finally: - remaining: Final = _HELD_ATTEMPTS.get(slot, 1) - 1 - if remaining > 0: - _HELD_ATTEMPTS[slot] = remaining - else: - _HELD_ATTEMPTS.pop(slot, None) + yield LoginAttempt(throttle=self, username=username, block=block, slot=slot) async def _active_block(self, keys: _Keys) -> Block | None: local: Final = self._local_block_ttls(keys) @@ -391,6 +373,7 @@ class LoginAttempt: throttle: LoginThrottle username: str block: Block | None + slot: str | None = None async def succeeded(self) -> None: if not self.throttle.enabled: @@ -400,9 +383,8 @@ class LoginAttempt: async def failed(self) -> None: if not self.throttle.enabled: return - if self.block is not None: - await _sleep(BLOCKED_ATTEMPT_HOLD_SECONDS) - self.throttle.refuse(max(self.block.retry_after - BLOCKED_ATTEMPT_HOLD_SECONDS, 1)) + if self.block is not None and self.slot is not None: + await self._hold_then_refuse(self.block, self.slot) user_block, source_block = await self.throttle.record_failure(self.username) if user_block == 0 and source_block == 0: return @@ -413,3 +395,26 @@ class LoginAttempt: self.username, self.throttle.client_ip, ) + + async def _hold_then_refuse(self, block: Block, slot: str) -> NoReturn: + held: Final = _HELD_ATTEMPTS.get(slot, 0) + if held >= MAX_HELD_ATTEMPTS_PER_KEY: + verbose_proxy_logger.warning( + "Admin UI sign-in refused at once: %s wrong attempts already held for a blocked %s; " + "username=%r source=%s", + held, + block.scope, + self.username, + self.throttle.client_ip, + ) + self.throttle.refuse(block.retry_after) + _HELD_ATTEMPTS[slot] = held + 1 + try: + await _sleep(BLOCKED_ATTEMPT_HOLD_SECONDS) + finally: + remaining: Final = _HELD_ATTEMPTS.get(slot, 1) - 1 + if remaining > 0: + _HELD_ATTEMPTS[slot] = remaining + else: + _HELD_ATTEMPTS.pop(slot, None) + self.throttle.refuse(max(block.retry_after - BLOCKED_ATTEMPT_HOLD_SECONDS, 1)) diff --git a/tests/test_litellm/proxy/auth/test_login_utils.py b/tests/test_litellm/proxy/auth/test_login_utils.py index df5bd29f8ff..df35e916aa0 100644 --- a/tests/test_litellm/proxy/auth/test_login_utils.py +++ b/tests/test_litellm/proxy/auth/test_login_utils.py @@ -1211,7 +1211,7 @@ async def test_held_attempts_from_one_blocked_key_are_capped(monkeypatch): with pytest.raises(ProxyException) as over_cap: await _guess(throttle) assert over_cap.value.code == "429" - assert over_cap.value.headers.get("Retry-After") == "30" + assert over_cap.value.headers.get("Retry-After") == "300", "refused at once, for the whole block" assert await _fail(throttle, username="someone-else@corp.com") == "401", "other keys are not affected" finally: release.set() @@ -1222,6 +1222,46 @@ async def test_held_attempts_from_one_blocked_key_are_capped(monkeypatch): assert lt._HELD_ATTEMPTS.get(slot) is None, "the slots are released once the held attempts answer" +@pytest.mark.asyncio +async def test_a_full_hold_pool_still_lets_the_right_password_in(monkeypatch): + """Five parked wrong guesses from the office must not turn the soft block into a lockout for the real user.""" + import asyncio + + from litellm.proxy.auth import login_throttle as lt + from litellm.proxy.auth.login_throttle import MAX_HELD_ATTEMPTS_PER_KEY + + monkeypatch.setenv("UI_USERNAME", "admin") + monkeypatch.setenv("UI_PASSWORD", "right") + release = asyncio.Event() + + async def _park(_seconds: float) -> None: + await release.wait() + + monkeypatch.setattr(lt, "_sleep", _park) + throttle = _throttle(user_limit=1, source_limit=3, client_ip="203.0.113.46") + assert [await _fail(throttle, username=f"spray-{i}@corp.com") for i in range(4)] == ["401"] * 4 + source_slot = throttle._keys("known@example.com").source_block + assert throttle._local_block_ttl(source_slot) > 0, "the source is blocked" + + held = [ + asyncio.create_task(_guess(throttle, username="known@example.com")) for _ in range(MAX_HELD_ATTEMPTS_PER_KEY) + ] + for _ in range(1000): + if lt._HELD_ATTEMPTS.get(source_slot) == MAX_HELD_ATTEMPTS_PER_KEY: + break + await asyncio.sleep(0) + assert lt._HELD_ATTEMPTS == {source_slot: MAX_HELD_ATTEMPTS_PER_KEY} + + try: + signed_in = await _db_login(throttle, "known@example.com", "right", correct=True) + assert signed_in.user_id == "u-1" + finally: + release.set() + for task in held: + with pytest.raises(ProxyException): + await task + + @pytest.mark.asyncio async def test_a_blocked_source_shares_one_held_slot_pool_across_its_blocked_usernames(monkeypatch): """Once the source is blocked, a pair block for a username must not hand that username its own five slots.""" @@ -1258,7 +1298,7 @@ async def test_a_blocked_source_shares_one_held_slot_pool_across_its_blocked_use with pytest.raises(ProxyException) as over_cap: await _guess(throttle, username=name) assert over_cap.value.code == "429" - assert over_cap.value.headers.get("Retry-After") == "30" + assert over_cap.value.headers.get("Retry-After") == "300" finally: release.set() for task in held: