对FastApi的API进行限速

Python研究所

发布于 2022-06-17 08:12:12

1.1K0

发布于 2022-06-17 08:12:12

文章被收录于专栏：大飞的部落阁大飞的部落阁

前言

今天啊，我的朋友小绿跟我说，那天晚上有个人好快好快！我急忙问他，怎么回事？有多快！他说，他的资源分发接口被人以很快的速度攻击了。速度达到了 10000qps，这对于他那 0.1C125M 的机器来说简直是小车拉大牛，真费车。

不忍心他承受那么快的速度，出于关心，我决定为他的接口加上限速。这样不但可以承受高压，还不影响正常服务用户，简直舒服。

关于 walrus

walrus 是一个 Redis 的操作库（对 redis-py 的二次封装），使用它可以简单实现数据缓存、全文搜索、频率限制等功能。今天我们就要使用 walrus 的频率限制功能来实现接口的频率限制。这对那些低频接口，一次性接口非常有用。

呆猫

安装 walrus

pip install walrus

安装 Redis

官方地址：https://redis.io/

以下为 windows 版的 Redis 包下载地址

链接：https://pan.baidu.com/s/1ZRzMaS5eqCPaotX5DljY3w&shfl=sharepset 提取码：spf0

本文使用 windows 作为演示，Linux 同理。

下载完成后解压后
cd 到 redis 文件夹
在CMD中执行 redis-server.exe redis.windows.conf启动 redis 服务
然后重新打开一个CMD使用 redis-cli.exe -h 127.0.0.1 -p 6379 登录 reids 服务

未做限速的后端代码

from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
import uvicorn

app = FastAPI()

@app.get('/')
def index():
    return {'status_code': True}


@app.get('/only2')
def api1(request: Request):
    return {'status_code': True, 'data': '不限速！'}

if __name__=='__main__':
    uvicorn.run(app='demo1:app',host='localhost',port=1213,reload=True)

以上代码是不限速的，对于only2接口的访问可以无限，存在被并发攻击的风险。对于分发资源的接口来说非常不好。

限速后的后端代码

from walrus import Database, RateLimitException
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
import uvicorn

#db = Database(host='localhost',port=6379)
# 默认连接loalhost:6379,所以记得启动本地redis服务哦
db = Database()
rate = db.rate_limit('speedlimit', limit=2, per=60) # 每分钟只能调用2次

app = FastAPI()


@app.exception_handler(RateLimitException)
def parse_rate_litmit_exception(request: Request, exc: RateLimitException):
    msg = {'status_code': False, 'msg': f'别那么快，好不好？我已经记住了你的IP: {request.client.host}'}
    return JSONResponse(status_code=412, content=msg)

@app.get('/')
def index():
    return {'status_code': True}


@app.get('/only2')
@rate.rate_limited(lambda request: request.client.host)
def api2(request: Request):
    return {'status_code': True, 'data': '这个速度可以！'}

if __name__=='__main__':
    uvicorn.run(app='demo:app',host='localhost',port=1212,reload=True)