mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
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:
parent
438b4e6a3f
commit
0a8423d77b
2 changed files with 69 additions and 24 deletions
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue