Flask-Limiter与认证：实现未认证用户优先返回401而非429的策略

from flask import Flask, jsonify, request
from flask_limiter import Limiter
from flask_limiter.util import get_remote_address
from functools import wraps

app = Flask(__name__)

# 初始化Flask-Limiter
# 配置默认限速规则，并使用内存存储。
# 在生产环境中，推荐使用Redis等持久化存储。
limiter = Limiter(
    app=app,
    key_func=get_remote_address, # 默认根据客户端IP进行限速
    default_limits=["1 per day", "1 per hour"],
    storage_uri="memory://",
)

# 模拟用户认证函数
def is_authenticated():
    """
    此函数模拟实际的用户认证逻辑。
    在真实应用中，这里会根据请求头（如Authorization）、
    会话（session）或Cookie等检查用户的认证凭证。
    为演示目的，我们假设用户始终未认证。
    """
    # 实际场景中，可能根据 token 或 session 判断
    # return request.headers.get('Authorization') == 'Bearer my_secret_token'
    return False # 默认返回False，模拟未认证用户

# 自定义认证装饰器
def authenticated_request(f):
    """
    一个用于保护路由的认证装饰器。
    在当前设计中，由于before_request已处理未认证情况，
    此装饰器更多作为路由层面的额外安全层或特定认证逻辑。
    """
    @wraps(f)
    def decorated_function(*args, **kwargs):
        if not is_authenticated():
            # 如果before_request已返回401，这段代码通常不会被执行。
            # 但作为防御性编程，保留此检查。
            print('Decorator: User not authenticated.')
            return jsonify({"message": "Unauthorized from decorator"}), 401
        print('Decorator: User authenticated.')
        return f(*args, **kwargs)
    return decorated_function

# before_request 钩子：优先处理认证
@app.before_request
def check_auth_and_rate_limit():
    """
    在每个请求处理之前执行。
    首先检查用户认证状态。如果未认证，立即返回401。
    如果已认证，则继续进行限速检查。
    """
    print(f"[{request.path}] Checking authentication and rate limit...")
    if not is_authenticated():
        print(f"[{request.path}] User not authenticated. Returning 401.")
        # 如果用户未认证，立即返回401响应，中断后续请求处理。
        return jsonify({"message": "Unauthorized"}), 401
    else:
        print(f"[{request.path}] User is authenticated. Checking rate limit...")
        # 用户已认证，现在进行限速检查。
        # limiter.check() 会返回一个元组 (limit_hit, limit_info)
        resp = limiter.check()
        if resp and resp[1]: # 如果限速触发 (resp[1] 为 True)
            print(f"[{request.path}] Rate limit exceeded for authenticated user. Returning 429.")
            return jsonify({"message": "Rate limit exceeded"}), 429
    print(f"[{request.path}] Authentication passed and rate limit not hit. Continuing request.")
    # 如果认证通过且未触发限速，则请求继续正常处理。

# 示例路由
@app.route('/example')
@authenticated_request # 路由层面的认证装饰器
def example_route():
    """
    一个受认证保护的示例路由。
    """
    return jsonify({"message": "This is an example route (authenticated access)"})

if __name__ == '__main__':
    # 运行Flask应用，debug模式下方便观察输出
    app.run(debug=True, port=5000)