fastapi-limiter

简介

FastAPI-Limiter 是一个使用 Lua 脚本为 fastapi 路由提供速率限制的工具。

要求

redis

安装

只需从 pypi 安装

> pip install fastapi-limiter

快速开始

FastAPI-Limiter 使用简单，只提供了一个依赖 RateLimiter，以下示例允许在路由 / 上每 5 秒请求 2 次。

import redis.asyncio as redis
import uvicorn
from contextlib import asynccontextmanager
from fastapi import Depends, FastAPI

from fastapi_limiter import FastAPILimiter
from fastapi_limiter.depends import RateLimiter

@asynccontextmanager
async def lifespan(_: FastAPI):
    redis_connection = redis.from_url("redis://localhost:6379", encoding="utf8")
    await FastAPILimiter.init(redis_connection)
    yield
    await FastAPILimiter.close()

app = FastAPI(lifespan=lifespan)

@app.get("/", dependencies=[Depends(RateLimiter(times=2, seconds=5))])
async def index():
    return {"msg": "Hello World"}


if __name__ == "__main__":
    uvicorn.run("main:app", debug=True, reload=True)

用法

FastAPILimiter.init 中有一些配置。

redis

aioredis 的 redis 实例。

prefix

Redis 键的前缀。

identifier

路由限制的标识符，默认为 ip，你可以重写它，例如使用 userid 等。

async def default_identifier(request: Request):
    forwarded = request.headers.get("X-Forwarded-For")
    if forwarded:
        return forwarded.split(",")[0]
    return request.client.host + ":" + request.scope["path"]

callback

访问被禁止时的回调，默认是抛出状态码为 429 的 HTTPException。

async def default_callback(request: Request, response: Response, pexpire: int):
    """
    请求过多时的默认回调
    :param request:
    :param pexpire: 剩余的毫秒数
    :param response:
    :return:
    """
    expire = ceil(pexpire / 1000)

    raise HTTPException(
        HTTP_429_TOO_MANY_REQUESTS, "Too Many Requests", headers={"Retry-After": str(expire)}
    )

多个限制器

你可以在一个路由中使用多个限制器。

@app.get(
    "/multiple",
    dependencies=[
        Depends(RateLimiter(times=1, seconds=5)),
        Depends(RateLimiter(times=2, seconds=15)),
    ],
)
async def multiple():
    return {"msg": "Hello World"}

注意你应该注意依赖的顺序，将 seconds/times 的结果较小的放在前面。

WebSocket 中的速率限制

虽然上述示例适用于 REST 请求，但 FastAPI 也允许轻松使用 WebSocket，这需要稍微不同的方法。

由于 WebSocket 可能会长时间存在，你可能希望对通过 socket 发送的数据进行速率限制。

你可以通过在 WebSocket 处理程序的主体中进行速率限制来实现这一点：

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    ratelimit = WebSocketRateLimiter(times=1, seconds=5)
    while True:
        try:
            data = await websocket.receive_text()
            await ratelimit(websocket, context_key=data)  # 注意：context_key 是可选的
            await websocket.send_text(f"Hello, world")
        except WebSocketRateLimitException:  # 超出速率限制时抛出
            await websocket.send_text(f"Hello again")

Lua 脚本

使用的 Lua 脚本。

local key = KEYS[1]
local limit = tonumber(ARGV[1])
local expire_time = ARGV[2]

local current = tonumber(redis.call('get', key) or "0")
if current > 0 then
    if current + 1 > limit then
        return redis.call("PTTL", key)
    else
        redis.call("INCR", key)
        return 0
    end
else
    redis.call("SET", key, 1, "px", expire_time)
    return 0
end