Flask Limiter - Using more than one key_func for a single endpoint

user_5347176 0
Viewed 2191

I'm working on building up an API endpoint w/ flask and flask_limiter, and I'm running into an issue. My goal is to rate limit on BOTH the client IP address, as well as an API key that the client provides (accessible through flask's "request" module), on a single endpoint.

The IP address limiter is important because I'd like to protect against brute force attacks, whereas the API key limiter is simply for business reasons. I've built up both rate limiters, and independently they work (i.e. I can limit on the IP address OR the provided API key), but I haven't been able to get both running on the endpoint at the same time. An example of what I've tried out so far:

from flask import Flask, request
from flask_limiter import Limiter
from flask_limiter.util import get_remote_address

app = Flask(__name__)
ip_limiter = Limiter(app, key_func=get_remote_address)
get_api_key = lambda : request.args.get('apikey')
api_key_limiter = Limiter(app, key_func=get_api_key)

@app.route('/theEndpoint')
@ip_limiter.limit('some_limit_here')
@api_key_limiter.limit('other_limit_here')
def theEndpointFunction(...):
    .......

It would be nice to implement all rate limits in a single package (i.e. flask_limiter), but thus far the only way I can think to implement rate limits on multiple keys is to switch over to something like redis.

Any clues on a way to implement multi-key rate-limits with flask_limiter?

flask api-design python
2 Answers

In case anyone is interested, I've decided to switch over to Redis to implement my dual-key rate limit. The solution goes something like this:

class RateLimit(object):
    def __init__(self, key, max_requests, seconds):
        self.reset = int(time.time()) + seconds
        self.key = key
        self.max_requests = max_requests
        self.seconds = seconds
        p = redis.pipeline()
        p.incr(self.key)
        p.expireat(self.key, self.reset)
        self.current = min(p.execute()[0], max_requests)
    remaining = property(lambda x: x.max_requests - x.current)
    over_limit = property(lambda x: x.current >= x.max_requests)


def get_view_rate_limit():
    return getattr(g, '_view_rate_limit', None)

def over_limit(limit):
    #formatting a JSON response
    response = {"result": "Max number of requests exceeded", 
                "status": False}
    response = jsonify(response)
    response.status_code = 400
    return response

def ratelimit(max_requests, seconds, key_func, over_limit=over_limit):
    def decorator(f):
        def rate_limited(*args, **kwargs):
            key = key_func()
            rlimit = RateLimit(key, max_requests, seconds)
            g._view_rate_limit = rlimit
            if over_limit is not None and rlimit.over_limit:
                return over_limit(rlimit)
            return f(*args, **kwargs)
        return update_wrapper(rate_limited, f)
    return decorator


RATE_LIMIT_FUNCS = {'apikey': lambda: request.args.get('apikey'),
                    'ip': lambda: get_remote_address}

and then decorating the flask endpoint:

@app.route('/theEndpoint')
@ratelimit(max_requests=5, seconds=300, key_func=RATE_LIMIT_FUNCS['apikey'])
@ratelimit(max_requests=15, seconds=300, key_func=RATE_LIMIT_FUNCS['ip'])
def theEndpointFunction():
    ....
user_53471760

Flask limiter has since provided support for different kinds of storage:

Either through config/env values stored in RATELIMIT_STORAGE_URI or passed when initializing the Limiter object through the storage_uri kwarg. The value expected is the storage type and the Host to the storage system, if needed.

https://flask-limiter.readthedocs.io/en/stable/configuration.html#RATELIMIT_STORAGE_URI

As the documentation notes, it supports

  • In memory (by default) with memory://
  • Memcached: memcached://host:port
  • MongoDB: mongodb://host:port
  • Redis: redis://host:port

example

ip_limiter = Limiter(
    app, 
    key_func=get_remote_address,
    storage_uri='redis://127.1.0.1:6379'
)
user_81330390
Related