Django REST framework之訪問頻率控制/節流

1. 訪問頻率控制基本實現

可以根據 ip地址 來對用戶訪問頻率進行限制，所以我們可以自定這樣的訪問頻率控制的類：

還可以通過具有唯一標識意義的用戶名或者用戶ID等！

throttle.py:

import time

VISIT_RECORD = {}


class VisitThrottle:
    def allow_request(self, request, view):
        """
        設定10s之內只能訪問10次
        :param request:
        :param view:
        :return: True or False
        返回值爲True表示可以訪問;返回值爲False或None表示訪問頻率太高被限制
        """
        # 獲取用戶的ip地址,當前request(封裝)中有的就取當前的request，如果沒有就到_request中取
        remote_addr = request._request.META.get('REMOTE_ADDR')
        print(remote_addr)  # 127.0.0.1
        ctime = time.time()
        if remote_addr not in VISIT_RECORD:
            VISIT_RECORD[remote_addr] = [ctime]
            return True
        history = VISIT_RECORD.get(remote_addr)
        while history and history[-1] < ctime - 10:
            history.pop(-1)
        if len(history) < 10:
            history.insert(0, ctime)
            return True
        return False  # 寫不寫都可以,如果執行到這裏說明不可訪問的，返回False或None都可以表示不可以訪問

    def wait(self, *args, **kwargs):
        pass

views.py:

from rest_framework.views import APIView
from .models import User, UserToken
from django.http import JsonResponse
from utils.md5 import md5
from app.throttle import VisitThrottle


class AuthView(APIView):
    authentication_classes = []
    permission_classes = []
    throttle_classes = [VisitThrottle, ]

    def post(self, request, *args, **kwargs):
        ret = {'code': 1000, 'msg': None}
        try:
            # 需要以form-data的方式提交
            name = request._request.POST.get('name')
            pwd = request._request.POST.get('pwd')
            instance = User.objects.filter(name=name, pwd=pwd).first()  # User object (1)，
            print(type(instance))  # <class 'app.models.User'>，加不加all()結果一樣
            print(instance)  # User object (1)，加不加all()結果一樣
            if not instance:
                ret['code'] = 1001
                ret['msg'] = '用戶名或密碼錯誤'
            else:
                token = md5(name=name)
                UserToken.objects.update_or_create(user=instance, defaults={'token': token})
                ret['token'] = token
        except Exception as e:
            ret['code'] = 1001
            ret['msg'] = '請求異常'
        return JsonResponse(ret)

10s內只能發送10次請求，超出10次請求，則會返回。請求被限制：

當請求被拒絕後，系統還可以返回 距離下次可以請求的時間，可以通過wait方法來設置：

throttle.py:

import time
# 可以放到緩存中，默認Django REST framework中就是放在緩存中
VISIT_RECORD = {}


class VisitThrottle:
    def __init__(self):
        self.history = None
        self.ctime = None

    def allow_request(self, request, view):
        """
        設定10s之內只能訪問10次
        :param request:
        :param view:
        :return: True or False
        返回值爲True表示可以訪問;返回值爲False或None表示訪問頻率太高被限制
        """
        # 獲取用戶的ip地址,當前request(封裝)中有的就取當前的request，如果沒有就到_request中取
        remote_addr = request._request.META.get('REMOTE_ADDR')
        print(remote_addr)  # 127.0.0.1
        self.ctime = time.time()  # 1593151531.1494734
        print(self.ctime)
        if remote_addr not in VISIT_RECORD:
            VISIT_RECORD[remote_addr] = [self.ctime]
            return True
        self.history = VISIT_RECORD.get(remote_addr)
        while self.history and self.history[-1] < self.ctime - 10:
            self.history.pop(-1)
        if len(self.history) < 10:
            self.history.insert(0, self.ctime)
            return True
        return False  # 寫不寫都可以,如果執行到這裏說明不可訪問的，返回False或None都可以表示不可以訪問

    def wait(self, *args, **kwargs):
        """
        設置距離下次可以請求的時間
        :param args:
        :param kwargs:
        :return:默認返回None，表示使用默認
        """
        """
        當請求被拒絕，會執行wait方法
        """
        ctime = time.time()
        print(ctime)
        print(self.history[-1])
        """
        return False
        {
          "detail": "Request was throttled. Expected available in 0 seconds."
        }
        """
        return 60 - (ctime - self.history[-1])

2. 訪問頻率控制源碼流程

請求過來先走dispatch方法：

執行當前self.initialize_request方法：

執行self.check_throttles方法：

執行self.get_throttles方法，遍歷訪問控制類對象列表：

通過列表生成式生成訪問控制類對象列表：

默認的訪問控制類是從配置文件中獲取的，也是全局的訪問控制類：


如果訪問被拒絕，即執行allow_request方法返回False或者None，則執行throttle.wait()方法：

3. 訪問頻率全局配置

同認證和權限一樣在配置文件中進行配置：

settings.py:

REST_FRAMEWORK = {
    'DEFAULT_THROTTLE_CLASSES': ['app.throttle.VisitThrottle',]
}

同樣某些請求可以不使用訪問頻率控制：

class AuthView(APIView):
    authentication_classes = []
    permission_classes = []
    # 訪問頻率控制
    throttle_classes = []

    def post(self, request, *args, **kwargs):
        pass

4. 內置訪問頻率控制類

Django REST framewor中內置了五個訪問頻率控制類，所有的訪問權限控制類應該繼承BaseThrottle並重寫其中的allow_request(self, request, view)方法：


內置的訪問控制類SimpleRateThrottle其實幫助我們實現了基於IP的訪問權限控制，我們看它內部是怎麼實現的：


所以之前自定義的類實現的功能完全可以繼承SimpleRateThrottle來實現，首先自定義類繼承SimpleRateThrottle，並設置scope：

from rest_framework.throttling import SimpleRateThrottle

class VisitThrottle(SimpleRateThrottle):
    # scope被當作key使用的，根據它到配置文件中取值
    scope = 'erics'

在配置文件中添加DEFAULT_THROTTLE_RATES：

REST_FRAMEWORK = {
    'DEFAULT_THROTTLE_RATES': {
        # 每分鐘訪問3次
        'erics': '3/m'
    }
}

SimpleRateThrottle類中的邏輯是首先根據scope從配置文件中獲取訪問頻率配置 3/m：

然後執行將獲取到的rate也就是scope傳入parse_rate函數進行解析：


__init__構造函數執行完成之後，執行allow_request方法：

獲取key，把ip當作key：

class VisitThrottle(SimpleRateThrottle):
    # scope被當作key使用的，根據它到配置文件中取值
    scope = 'erics'

    def get_cache_key(self, request, view):
        """獲取key"""
        return self.get_ident()

接下來到緩存獲取所有記錄：



之前自定義的時間，wait方法也已經做好，我們不需要去寫：

最終我們只需要幾行代碼就可以完成對匿名用戶的訪問頻率的控制：

from rest_framework.throttling import SimpleRateThrottle

class VisitThrottle(SimpleRateThrottle):
    # scope被當作key使用的，根據它到配置文件中取值
    scope = 'erics'

    def get_cache_key(self, request, view):
        """獲取key"""
        return self.get_ident(request)

REST_FRAMEWORK = {
    'DEFAULT_THROTTLE_CLASSES': ['app.throttle.VisitThrottle', ],
    'DEFAULT_THROTTLE_RATES': {
        # 每分鐘訪問5次,/右邊只要首字母是m就可以了
        'erics': '5/m'
    }
}

匿名用戶每分鐘只能訪問5次：

上面是對匿名用戶的訪問登錄控制，也可以對登錄用戶做訪問頻率控制。只需要根據用戶名來做頻率控制條件：

from rest_framework.throttling import SimpleRateThrottle

class UserVisitThrottle(SimpleRateThrottle):
    # scope被當作key使用的，根據它到配置文件中取值
    scope = 'user_erics'

    def get_cache_key(self, request, view):
        """獲取key"""
        print(request)  # <rest_framework.request.Request object at 0x7f155c6b86d0>
        # 用戶認證成功之後就會有request.user
        print(request.user)  # User object (1)
        print(request.user.name)  # thanlon
        return request.user.name

REST_FRAMEWORK = {
    # 'DEFAULT_THROTTLE_CLASSES': ['app.throttle.VisitThrottle', ],
    # 登錄的用戶使用根據用戶名來做頻率限制，匿名用戶使用IP來做頻率限制。這裏全局配置了登錄用戶的，匿名用戶需要可以單獨設置通過ip來控制
    'DEFAULT_THROTTLE_CLASSES': ['app.throttle.UserVisitThrottle', ],
    'DEFAULT_THROTTLE_RATES': {
        # 每分鐘訪問5次,/右邊只要首字母是m就可以了
        'erics': '5/m',
        'user_erics': '10/m'
    }
}

登錄的用戶每分鐘只能訪問10次：