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>
This commit is contained in:
yucheng 2026-09-17 00:28:39 +00:00
parent 438b4e6a3f
commit 0a8423d77b
2 changed files with 69 additions and 24 deletions

View file

@ -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))

View file

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