ratelimit with redis and zrange

throttling.py

from __future__ import annotations

from django.conf import settings
from django.utils import timezone

from rest_framework import throttling

from django_redis import get_redis_connection
from django_redis.client import DefaultClient
from ipware import get_client_ip
from redis import Redis

cache: DefaultClient


class ApiThrottle(throttling.BaseThrottle):
    @staticmethod
    def _ban_address_cache_key(address: str, endpoint: str) -> str:
        return f'ban:address:{address}:{endpoint}'

    @staticmethod
    def _ban_ssid_cache_key(ssid: str, endpoint: str) -> str:
        return f'ban:ssid:{ssid}:{endpoint}'

    @staticmethod
    def _request_log_address_cache_key(remote_addr: str, endpoint: str) -> str:
        return f'request:address:{remote_addr}:{endpoint}'

    @staticmethod
    def _request_log_ssid_cache_key(ssid: str, endpoint: str) -> str:
        return f'request:ssid:{ssid}:{endpoint}'

    def allow_request(self, request, view):
        """Реализация ratelimit со скользящим окном по логам в redis.

        Держим в логе только нужные запросы (zremrangebyscore),
        Быстро считаем их количество (zcard и zcount).
        Средствами redis чистим старые логи (expire).
        Все комманды в Lua для уменьшения запросов к серверу Redis
        и исключения одновременного доступа к одному ключу.
        """

        if not hasattr(request, 'session'):
            return False

        current_dt = timezone.now()
        if hasattr(view, 'ban_settings') and view.action in view.ban_settings:
            ban_settings = view.ban_settings[view.action]
        else:
            ban_settings = settings.BAN_REQUEST_SETTINGS

        request.session.setdefault('ban', 0)

        if request.session.session_key is None:
            request.session.save()

        ssid = request.session.session_key
        remote_addr, routable = get_client_ip(request)

        if remote_addr is None:
            return False

        redis_con: Redis = get_redis_connection()

        endpoint = f'{view.basename}_{view.action}'

        lua = """
            local timestamp = ARGV[1]
            local session = ARGV[2]
            local min_timeout = ARGV[3]
            local max_timeout = ARGV[4]
            local min_count = tonumber(ARGV[5])
            local max_count = tonumber(ARGV[6])
            local max_session_for_ip = tonumber(ARGV[7])
            local ban_timeout = tonumber(ARGV[8])

            redis.call('ZREMRANGEBYSCORE', KEYS[1], 0, timestamp - max_timeout * 60)
            redis.call('ZADD', KEYS[1], timestamp, timestamp)
            redis.call('EXPIRE', KEYS[1], max_timeout * 60)

            redis.call('ZREMRANGEBYSCORE', KEYS[2], 0, timestamp - 60 * 60)
            redis.call('ZADD', KEYS[2], timestamp, session)
            redis.call('EXPIRE', KEYS[2], 60 * 60)

            if tonumber(redis.call('EXISTS', KEYS[3], KEYS[4])) > 0 then
                return 1
            end

            if (
                (max_count and tonumber(redis.call('ZCARD', KEYS[1])) > max_count) or
                tonumber(redis.call('ZCOUNT', KEYS[1], timestamp - min_timeout * 60, '+inf')) > min_count
            ) then
                redis.call('SET', KEYS[3], 1, 'EX', ban_timeout * 60)
                return 1
            end

            if tonumber(redis.call('ZCARD', KEYS[2])) > max_session_for_ip then
                redis.call('SET', KEYS[4], 1, 'EX', ban_timeout * 60)
                return 1
            else
                return 0
            end
        """
        script = redis_con.register_script(lua)
        if script(
            keys=[
                self._request_log_ssid_cache_key(ssid, endpoint),
                self._request_log_address_cache_key(remote_addr, endpoint),
                self._ban_ssid_cache_key(ssid, endpoint),
                self._ban_address_cache_key(remote_addr, endpoint),
            ],
            args=[
                current_dt.timestamp(),
                ssid,
                ban_settings['min']['minutes'],
                # Если max нет, то подменяем на min и обрезаем очередь по нему
                ban_settings.get('max', {}).get('minutes')
                or ban_settings['min']['minutes'],
                ban_settings['min']['count'],
                ban_settings.get('max', {}).get('count', 'nil'),
                ban_settings['max_session_amount_for_ip'],
                ban_settings['ban_minutes'],
            ],
        ):
            return False

        return True