Implement a python decorator to limit flask request frequency by path and IP with the help of redis

I have a flask application. Some of its APIs are very heavy and I want to limit people's requests to them.

For example, I want that people with the same IP can only access the payment API for one time in 10 seconds.

Implement the decorator

The general idea is:

  1. Once a user makes a request to my protectd API, we mark it in redis with expiration.
  2. The next time when the same user makes a request, we check if there's a mark in redis.
    1. If the mark exists, we forbid this new request.
    2. If the mark expired, we handle it as normal.

Talk is cheap, here's the code:

from flask import Flask, abort, request
from functools import wraps
import redis

redis_conn = redis.Redis()
app = Flask(__name__)


def limit_request(interval=3):
    "Only allow one request in a specified interval (3 seconds by default)"
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            key = f'limit_request:{request.path}:{request.remote_addr}'
            val = redis_conn.get(key)
            if val is not None:
                # 429 Too Many Requests
                abort(429)
            else:
                # mark in redis and expire after interval seconds
                redis_conn.set(key, '1', ex=interval)
            rt = func(*args, **kwargs)
            return rt
        return wrapper
    return decorator


# an IP can only request "/pay" one time in 10 seconds
@app.route('/pay')
@limit_request(10)
def pay():
    return 'pay'


if __name__ == '__main__':
    app.run()

Note:

  • We put the mark in redis, so that multiple flask process (maybe on multiple machines) can all access the mark.
  • The redis key f'limit_request:{request.path}:{request.remote_addr}' contains prefix, request path and request IP. Therefore, different APIs and different IPs do not interfere with each other.

What if I don't want distinction between users (IPs)

Just remove request.remote_addr from the redis key:

# before
key = f'limit_request:{request.path}:{request.remote_addr}'
# after
key = f'limit_request:{request.path}'

request.remote_addr is not accurate?

If you put your flask server behind a proxy like Nginx, then request.remote_addr may be the IP of Nginx.

To fix it, refer to Tell Flask it is Behind a Proxy

Posted on 2022-06-17