Kris
ecd3c90e80
新增设备身份管理、认证限流、并发通道、Webhook路由、RBAC权限控制、SSE/轮询降级等全套网关通道功能，包含： 1. 设备身份生成与签名验证 2. 设备令牌认证与速率限制 3. 内存+数据库双重设备注册表 4. 并发通道限流管理 5. Webhook安全处理与路由 6. RBAC权限校验系统 7. OpenAI API兼容适配层 8. Tailscale认证支持 9. HTTP轮询降级机制
112 lines
3.4 KiB
Python
112 lines
3.4 KiB
Python
import logging
|
|
import threading
|
|
import time
|
|
from enum import StrEnum
|
|
from typing import NamedTuple
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
AUTH_FAIL_LIMIT = 10
|
|
AUTH_LOCKOUT_SECONDS = 300
|
|
AUTH_RATE_SCOPE_BUCKETS = 10_000
|
|
|
|
|
|
class AuthRateScope(StrEnum):
|
|
TOKEN = "token"
|
|
PASSWORD = "password"
|
|
DEVICE_TOKEN = "device_token"
|
|
BOOTSTRAP = "bootstrap"
|
|
|
|
|
|
class _ScopeState(NamedTuple):
|
|
failures: int
|
|
locked_until: float
|
|
|
|
|
|
class AuthRateLimiter:
|
|
def __init__(
|
|
self,
|
|
fail_limit: int = AUTH_FAIL_LIMIT,
|
|
lockout_seconds: float = AUTH_LOCKOUT_SECONDS,
|
|
max_buckets: int = AUTH_RATE_SCOPE_BUCKETS,
|
|
):
|
|
self._fail_limit = fail_limit
|
|
self._lockout_seconds = lockout_seconds
|
|
self._max_buckets = max_buckets
|
|
self._buckets: dict[tuple[str, AuthRateScope], _ScopeState] = {}
|
|
self._lock = threading.Lock()
|
|
|
|
def _cleanup_expired(self) -> None:
|
|
if len(self._buckets) <= self._max_buckets:
|
|
return
|
|
now = time.monotonic()
|
|
expired = [
|
|
key
|
|
for key, state in self._buckets.items()
|
|
if now >= state.locked_until and state.failures < self._fail_limit
|
|
]
|
|
for key in expired:
|
|
del self._buckets[key]
|
|
|
|
def record_failure(self, identifier: str, scope: AuthRateScope) -> bool:
|
|
with self._lock:
|
|
self._cleanup_expired()
|
|
key = (identifier, scope)
|
|
now = time.monotonic()
|
|
state = self._buckets.get(key)
|
|
|
|
if state is not None and now < state.locked_until:
|
|
return False
|
|
|
|
new_failures = (state.failures + 1) if state else 1
|
|
locked_until = now + self._lockout_seconds if new_failures >= self._fail_limit else 0.0
|
|
self._buckets[key] = _ScopeState(new_failures, locked_until)
|
|
|
|
if new_failures >= self._fail_limit:
|
|
logger.warning(
|
|
"Auth rate limit locked: identifier=%s scope=%s",
|
|
identifier,
|
|
scope.value,
|
|
)
|
|
|
|
if len(self._buckets) > self._max_buckets * 2:
|
|
self._force_evict(now)
|
|
|
|
return True
|
|
|
|
def reset(self, identifier: str, scope: AuthRateScope) -> None:
|
|
key = (identifier, scope)
|
|
with self._lock:
|
|
self._buckets.pop(key, None)
|
|
|
|
def is_locked(self, identifier: str, scope: AuthRateScope) -> bool:
|
|
key = (identifier, scope)
|
|
with self._lock:
|
|
state = self._buckets.get(key)
|
|
if state is None:
|
|
return False
|
|
return time.monotonic() < state.locked_until
|
|
|
|
def remaining_attempts(self, identifier: str, scope: AuthRateScope) -> int:
|
|
key = (identifier, scope)
|
|
with self._lock:
|
|
state = self._buckets.get(key)
|
|
if state is None:
|
|
return self._fail_limit
|
|
now = time.monotonic()
|
|
if state.locked_until > 0.0 and now >= state.locked_until:
|
|
return self._fail_limit
|
|
return max(0, self._fail_limit - state.failures)
|
|
|
|
def _force_evict(self, now: float) -> None:
|
|
by_expiry = sorted(
|
|
self._buckets.items(),
|
|
key=lambda item: item[1].locked_until,
|
|
)
|
|
to_remove = len(self._buckets) - self._max_buckets
|
|
for i in range(to_remove):
|
|
del self._buckets[by_expiry[i][0]]
|
|
|
|
|
|
auth_rate_limiter = AuthRateLimiter()
|