fix rate limit being global instead of per-IP
Build and publish bnuy api / build (push) Successful in 14m0s
Build and publish bnuy api / build (release) Successful in 13m46s

This commit is contained in:
2026-05-16 22:37:00 +02:00
parent 716295b7fd
commit b56b899595
+21 -2
View File
@@ -4,7 +4,7 @@ from fastapi import FastAPI, HTTPException, Request, UploadFile
from fastapi.responses import FileResponse from fastapi.responses import FileResponse
from fastapi import Depends, FastAPI from fastapi import Depends, FastAPI
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from pyrate_limiter import Duration, Limiter, Rate from pyrate_limiter import AbstractBucket, BucketFactory, Duration, InMemoryBucket, Limiter, MonotonicClock, Rate, RateItem
from fastapi_limiter.depends import RateLimiter from fastapi_limiter.depends import RateLimiter
from fastapi import Depends from fastapi import Depends
import asyncmy import asyncmy
@@ -60,8 +60,27 @@ async def connect_db(app: FastAPI):
app.state.pool.close() app.state.pool.close()
await app.state.pool.wait_closed() await app.state.pool.wait_closed()
class MultiBucketFactory(BucketFactory):
def __init__(self, rates, clock):
self.clock = clock
self.rates = rates
self.buckets = {}
def wrap_item(self, name: str, weight: int = 1) -> RateItem:
"""Time-stamping item, return a RateItem"""
now = self.clock.now()
return RateItem(name, now, weight=weight)
def get(self, item: RateItem) -> AbstractBucket:
if item.name not in self.buckets:
new_bucket = self.create(InMemoryBucket, self.rates)
self.buckets.update({item.name: new_bucket})
return self.buckets[item.name]
app = FastAPI(lifespan=connect_db) app = FastAPI(lifespan=connect_db)
limiter = Limiter(Rate(50, Duration.MINUTE)) rates = [Rate(50, Duration.MINUTE)]
limiter = Limiter(MultiBucketFactory(rates,MonotonicClock()))
async def create_tables(pool): async def create_tables(pool):
async with pool.acquire() as conn: async with pool.acquire() as conn: