Spaces:
Sleeping
Sleeping
File size: 6,066 Bytes
c91c7db e1104b3 c91c7db e1104b3 c91c7db | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 | from __future__ import annotations
import asyncio
import math
import time
from collections import defaultdict, deque
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from sqlalchemy.dialects.sqlite import insert as sqlite_insert
from app.security.context import AuthContext
from app.security.database import SecurityDatabase
from app.security.errors import RateLimitError
from app.security.models import RateLimit, utcnow
@dataclass(slots=True)
class RateLimitLease:
limiter: APIKeyRateLimiter
api_key_id: str
concurrent: bool
async def release(self) -> None:
if self.concurrent:
await self.limiter.release_job(self.api_key_id)
class APIKeyRateLimiter:
"""Low-latency per-key windows with durable aggregate counters for auditing."""
def __init__(self, database: SecurityDatabase) -> None:
self.database = database
self._lock = asyncio.Lock()
self._requests: dict[str, deque[float]] = defaultdict(deque)
self._uploads: dict[str, deque[float]] = defaultdict(deque)
self._concurrent: dict[str, int] = defaultdict(int)
self._daily_bytes: dict[tuple[str, str], int] = defaultdict(int)
self._categories: dict[tuple[str, str], deque[float]] = defaultdict(deque)
async def acquire(
self,
context: AuthContext,
*,
is_job: bool,
is_upload: bool,
uploaded_bytes: int,
) -> RateLimitLease:
now = time.time()
today = datetime.now(timezone.utc).date().isoformat()
retry_after = 0
async with self._lock:
requests = self._requests[context.api_key_id]
self._prune(requests, now - 60)
if len(requests) >= context.requests_per_minute:
retry_after = max(1, math.ceil(requests[0] + 60 - now))
uploads = self._uploads[context.api_key_id]
self._prune(uploads, now - 3600)
if not retry_after and is_upload and len(uploads) >= context.uploads_per_hour:
retry_after = max(1, math.ceil(uploads[0] + 3600 - now))
daily_key = (context.api_key_id, today)
daily_total = self._daily_bytes[daily_key]
if (
not retry_after
and uploaded_bytes
and daily_total + uploaded_bytes > context.processing_bytes_per_day
):
tomorrow = datetime.now(timezone.utc).replace(
hour=0, minute=0, second=0, microsecond=0
) + timedelta(days=1)
retry_after = max(1, int((tomorrow - datetime.now(timezone.utc)).total_seconds()))
if (
not retry_after
and is_job
and self._concurrent[context.api_key_id] >= context.concurrent_jobs
):
retry_after = 1
if retry_after:
raise RateLimitError(retry_after)
requests.append(now)
if is_upload:
uploads.append(now)
if uploaded_bytes:
self._daily_bytes[daily_key] += uploaded_bytes
if is_job:
self._concurrent[context.api_key_id] += 1
await self._record(context.api_key_id, "requests_minute", 1, 0, 60)
if is_upload:
await self._record(context.api_key_id, "uploads_hour", 1, 0, 3600)
if uploaded_bytes:
await self._record(
context.api_key_id, "processing_bytes_day", 0, uploaded_bytes, 86_400
)
return RateLimitLease(self, context.api_key_id, is_job)
async def release_job(self, api_key_id: str) -> None:
async with self._lock:
self._concurrent[api_key_id] = max(0, self._concurrent[api_key_id] - 1)
async def acquire_category(
self,
context: AuthContext,
category: str,
*,
limit: int,
window_seconds: int,
) -> None:
"""Reserve an independent social-operation bucket for a key.
Generic API limits still apply in middleware. These smaller buckets
prevent OAuth, publishing, scheduling, and analytics traffic from
starving each other when the social subsystem is enabled.
"""
now = time.time()
key = (context.api_key_id, category)
async with self._lock:
values = self._categories[key]
self._prune(values, now - window_seconds)
if len(values) >= limit:
raise RateLimitError(max(1, math.ceil(values[0] + window_seconds - now)))
values.append(now)
await self._record(context.api_key_id, category, 1, 0, window_seconds)
@staticmethod
def _prune(values: deque[float], cutoff: float) -> None:
while values and values[0] <= cutoff:
values.popleft()
async def _record(
self, api_key_id: str, bucket_type: str, count: int, units: int, seconds: int
) -> None:
now = datetime.now(timezone.utc)
epoch = int(now.timestamp())
bucket_start = datetime.fromtimestamp(epoch - (epoch % seconds), timezone.utc)
async with self.database.session() as session:
statement = sqlite_insert(RateLimit).values(
api_key_id=api_key_id,
bucket_type=bucket_type,
bucket_start=bucket_start,
count=count,
units=units,
updated_at=utcnow(),
)
statement = statement.on_conflict_do_update(
index_elements=[
RateLimit.api_key_id,
RateLimit.bucket_type,
RateLimit.bucket_start,
],
set_={
"count": RateLimit.count + statement.excluded.count,
"units": RateLimit.units + statement.excluded.units,
"updated_at": utcnow(),
},
)
await session.execute(statement)
await session.commit()
|