使用 pre-request 优化 Flask 入参校验

Posted 2021-09-06 Eric

在使用Flask实现Restful接口时，需要对请求参数进行校验以判断其是否符合特定的规则。本文将介绍如何通过 pre-request 库优化校验逻辑。

接口需求

假设我们需要实现一个收集用户填写的个人信息的接口，该接口的具体需求如下:

1. 接口路径: /user/info/new
2. 接口字段:

字段	类型	说明
userName	string	用户昵称，必填，2-20位字符串
gender	int	用户性别，必填，1-男 2-女
age	int	用户年龄，必填，18-60之间的整数
country	string	用户国籍，选填，默认为 中国, 字符串长度为>2

接口实现

如果不借助任何第三方工具实现上述接口的话，代码可能如下所示:

from flask import request, Flask


app = Flask(__name__)


@app.route("/user/info/new", methods=["POST"])
def user_info_handler():
    # 1.判断用户名称参数是否合法
    user_name = request.form.get("userName")
    if not user_name or not isinstance(user_name, str):
        return "请填写正确的用户名称"
    if len(user_name) < 2 or len(user_name) > 20:
        return "用户名称长度不正确"

    # 2.判断用户性别参数是否合法
    gender = request.form.get("gender")
    if not gender:
        return "请填写用户性别"
    try:
        gender = int(gender)
    except ValueError:
        return "用户性别格式不正确"
    if gender not in [1, 2]:
        return "用户性别参数必须在[1, 2]之间"

    # 3.判断用户年级参数是否合法
    age = request.form.get("age")
    if not age:
        return "请填写用户年龄"
    try:
        age = int(age)
    except ValueError:
        return "用户年龄格式不正确"
    if age < 18 or age > 60:
        return "年龄必须在18-60之间"

    # 4.判断国籍是否合法
    country = request.form.get("country", "中国")
    country = str(country)
    if len(country) < 2:
        return "国籍名称长度不合法"
    
    # TODO: 用户信息注册逻辑
    return "Success"


if __name__ == "__main__":
    app.run(port=8080)

正如上面的代码所示，为了保证存储到数据库中的数据符合设计要求，开发人员需要对入参做大量的校验工作，稍有不慎就有可能为系统留下漏洞隐患。

使用 pre-request 校验入参

同样是实现上述的接口需求，我们来看一下通过pre-request如何将大量重复的校验逻辑进行屏蔽。

from flask import Flask
from pre_request import pre, Rule


app = Flask(__name__)


rule = {
    "userName": Rule(type=str, required=True, gte=3, lte=20, dest="user_name"),
    "gender": Rule(type=int, required=True, enum=[1, 2]),
    "age": Rule(type=int, required=True, gte=18, lte=60),
    "country": Rule(type=str, required=False, gte=2, default="中国")
}


@app.route("/user/info/new", methods=["POST"])
def user_info_handler():
    params = pre.parse(rule=rule)
    # TODO: 用户信息注册逻辑
    return "Success"


if __name__ == "__main__":
    app.run(port=8080)

如上方所示，复杂的参数校验工作变成了校验规则的编写，pre.parse 函数会自动捕获请求参数并判断是否符合校验规则，不符合校验规则时会自动生成错误响应，用户不需要做任何处理。

pre-request 相关链接

• 代码地址: https://github.com/Eastwu5788...
• 文档地址: https://pre-request.readthedo...