Sanic 学习(二) – 开发一个限流扩展

今天继续以实践的方式来探索 Python 的 Web 开发框架：Sanic。在上一篇学习的文章中，我掌握了如何给 Sanic 开发扩展，并完成了一个 Session 功能扩展，今天我将尝试写一个 Sanic 限流功能扩展。

同 Session 一样，限流也是 Web 开发中不可或缺的实用功能。它能避免一些敏感页面，或是数据查询压力较大的页面被频繁访问。比如账号登录，我可以通过限流的方式让每个 IP，或是每个账号，在规定的时间内最多只能验证几次，从而降低账号被暴力破解的风险。

来看看我想要实现的功能应用代码：

from sanic import Sanic
from sanic.response import text

app = Sanic('zzxworld')

@app.get("
@RateLimit.set('3/minute')
async def home(request):
    return text("Welcome to zzx's World.")

跟一个正常的 Sanic 示例代码差不多，主要是 @RateLimit.set('3/minute') 这行。它意图是限制其下的接口方法每分钟只能访问 3 次。对于有 Python 开发经验的朋友，看到开头的 @ 符号就能明白，这会用到 Python 的「装饰器」语法。

因为前一篇文章中已经详述了如何开始 Sanic 的扩展开发，这篇就不再重复赘述。直接上功能实现的代码：

from sanic_ext import Extend, Extension
from limits import storage, strategies, parse
from functools import wraps

class RateLimit(Extension):
    """zzxworld 的限流扩展"""
    name="zzxRateLimit"

    _limiter = None

    def startup(self, bootstrap):
        """扩展入口"""
        # 创建限流数据存储
        limitStorage = storage.MemoryStorage()
        # 初始化限流功能实例
        self._limiter = strategies.MovingWindowRateLimiter(limitStorage)
        # 在每个请求中绑定限流功能实例
        self.app.request_middleware.appendleft(self.bindLimiter)

    def bindLimiter(self, request):
        """绑定限流实例对象到每个请求"""
        request.ctx.rateLimiter = self._limiter

    @staticmethod
    def set(rateConfig):
        """设置限流的装饰函数"""
        # 解析限流设置
        rateParam = parse(rateConfig)
        def decorator(func):
            @wraps(func)
            def invoke(request, *args, **kwargs):
                # 使用请求的函数名称加 IP 来判断请求频率是否超限
                if (request.ctx.rateLimiter.hit(
                    rateParam, func.__name__, request.client_ip)):
                    # 未超出限流设置，正常执行路由请求函数
                    return func(request, *args, **kwargs)
                else:
                    # 超出限流设置，输出 429 状态
                    return text('Too many requests.', 429)
            return invoke
        return decorator

Extend.register(RateLimit)

跟之前一样，核心逻辑都有注释，就不再做过多解释。唯一需要说明的是限流的核心处理部分我偷懒了，使用了 limits 这个 Python 包。它提供了比较完善的限流功能，比如限流数据的存储，除了上面用到的 MemoryStorage，还有支持 Redis 的 RedisStorage，而且自带三种限流策略。有这么好一个直接就能用的库，自然是拿来就用，什么都自己来写还是有点「肝」不动。