FastApi 中自定义装饰器

背景

FastAPI 开发 web 应用为了保护我们的接口安全,需要判断是否携带 token，如果没有的话就拒绝请求, 通常可以采用自定义装饰器或者拦截器中间件来实现，装饰器相对来说更灵活一些，本文介绍如何自定义一个装饰器来验证接口安全。 完成类似 Spring Security 的注解功能

创建可接受参数的自定义装饰器

创建一个接受参数的装饰器，通过定义一个外部函数来接收参数，内部函数来处理实际装饰逻辑。

# 自定义装饰器
def token_required(role: str):
    def decorator(func):
        @wraps(func)
        async def decorated_function(*args, **kwargs):
            request: Request = kwargs.get('request')
            if not request:
                raise HTTPException(status_code=400, detail="Bad request")

            token = request.headers.get('Authorization')
            if not token:
                raise HTTPException(status_code=401, detail="Token is missing")

            # 在这里添加对 token 和 role 的进一步验证逻辑
            # 假设 token 是 "expected_token" 并且 role 是 "admin"
            if token != "expected_token" or role != "admin":
                raise HTTPException(status_code=403, detail="Permission denied")

            return await func(*args, **kwargs)

        return decorated_function

    return decorator

使用装饰器

@app.get("/secure-data")
@token_required(role="admin")
async def secure_data(request: Request):
    return {"message": "This is a secure endpoint."}

启动服务

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8008)

测试

请求头没有携带 Authorization

{
  "detail": "Token is missing"
}

携带 错误 Authorization

{
  "detail": "Permission denied"
}

携带正确 Authorization

{
  "message": "This is a secure endpoint."
}