first commit

eadwinCode · eadwinCode · commit f11ad14669b0 · 2023-02-26T21:46:54.000+01:00
diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml
@@ -23,5 +23,3 @@ jobs:
           FLIT_USERNAME: ${{ secrets.FLIT_USERNAME }}
           FLIT_PASSWORD: ${{ secrets.FLIT_PASSWORD }}
         run: flit publish
-      - name: Deploy Documentation
-        run: make doc-deploy
diff --git a/.github/workflows/test_full.yml b/.github/workflows/test_full.yml
@@ -38,10 +38,10 @@ jobs:
       - name: Install Dependencies
         run: flit install --symlink
       - name: Black
-        run: black --check ellar tests
+        run: black --check ellar_throttler tests
       - name: isort
-        run: isort --check ellar tests
+        run: isort --check ellar_throttler tests
       - name: Flake8
-        run: flake8 ellar tests
+        run: flake8 ellar_throttler tests
       - name: mypy
-        run: mypy ellar
+        run: mypy ellar_throttler
diff --git a/LICENSE b/LICENSE
diff --git a/ellar_throttler/constants.py b/ellar_throttler/constants.py
@@ -0,0 +1,4 @@
+THROTTLER_LIMIT = "THROTTLER:LIMIT"
+THROTTLER_TTL = "THROTTLER:TTL"
+THROTTLER_OPTIONS = "THROTTLER:MODULE_OPTIONS"
+THROTTLER_SKIP = "THROTTLER:SKIP"
diff --git a/ellar_throttler/decorators.py b/ellar_throttler/decorators.py
@@ -0,0 +1,65 @@
+import typing as t
+
+from ellar.reflect import reflect
+
+from .constants import THROTTLER_LIMIT, THROTTLER_SKIP, THROTTLER_TTL
+
+
+def _set_throttler_metadata(target: t.Callable, limit: int, ttl: int) -> None:
+    reflect.define_metadata(THROTTLER_TTL, ttl, target)
+    reflect.define_metadata(THROTTLER_LIMIT, limit, target)
+
+
+def throttle(*, limit: int = 20, ttl: int = 60) -> t.Callable:
+    """
+    Adds metadata to the target which will be handled by the ThrottlerGuard to
+    handle incoming requests based on the given metadata.
+
+    Usage:
+    Use on controllers or individual route functions
+
+    @throttle(limit=20, ttl=300)
+    class ControllerSample:
+        @throttle(limit=20, ttl=300)
+        async def index(self):
+            ...
+
+    :param limit: Guard Type or Instance
+    :param ttl: Guard Type or Instance
+    :return: Callable
+    """
+
+    def decorator(func: t.Callable) -> t.Callable:
+        _set_throttler_metadata(func, limit=limit, ttl=ttl)
+        return func
+
+    return decorator
+
+
+def skip_throttle(skip: bool = True) -> t.Callable:
+    """
+    Adds metadata to the target which will be handled by the ThrottlerGuard
+    whether to skip throttling for this context.
+
+    Usage:
+    Use on controllers or individual route functions
+
+    @throttle(limit=20, ttl=300)
+    class ControllerSample:
+        @skip_throttle()
+        async def index(self):
+            ...
+
+        @skip_throttle(false)
+        async def create(self):
+            ...
+
+    :param skip:
+    :return:
+    """
+
+    def decorator(func: t.Callable) -> t.Callable:
+        reflect.define_metadata(THROTTLER_SKIP, skip, func)
+        return func
+
+    return decorator
diff --git a/ellar_throttler/exception.py b/ellar_throttler/exception.py
@@ -0,0 +1,28 @@
+import math
+import typing as t
+
+from ellar.core.exceptions import APIException
+from starlette import status
+
+
+class ThrottledException(APIException):
+    status_code = status.HTTP_429_TOO_MANY_REQUESTS
+    default_detail = "Request was throttled."
+    extra_detail_singular = "Expected available in {wait} second."
+    extra_detail_plural = "Expected available in {wait} seconds."
+
+    def __init__(self, wait: float = None, detail: t.Any = None) -> None:
+        if detail is None:
+            detail = self.default_detail
+        if wait is not None:
+            wait = math.ceil(wait)
+            detail = " ".join(
+                (
+                    detail,
+                    self.extra_detail_singular.format(wait=wait)
+                    if wait < 1
+                    else self.extra_detail_plural.format(wait=wait),
+                )
+            )
+        headers = {"Retry-After": "%d" % float(wait or 0.0)}
+        super().__init__(detail=detail, headers=headers)
diff --git a/ellar_throttler/interfaces/__init__.py b/ellar_throttler/interfaces/__init__.py
@@ -0,0 +1,5 @@
+from .throttler_storage import IThrottlerStorage
+
+__all__ = [
+    "IThrottlerStorage",
+]
diff --git a/ellar_throttler/interfaces/throttler_storage.py b/ellar_throttler/interfaces/throttler_storage.py
@@ -0,0 +1,26 @@
+import typing as t
+from abc import ABC, abstractmethod
+
+from ..throttler_storage_options import ThrottlerStorageOption
+from ..throttler_storage_record import ThrottlerStorageRecord
+
+
+class IThrottlerStorage(ABC):
+    @property
+    @abstractmethod
+    def storage(self) -> t.Dict[str, ThrottlerStorageOption]:
+        """
+        The internal storage with all the request records.
+        The key is a hashed key based on the current context and IP.
+        :return:
+        """
+
+    @abstractmethod
+    async def increment(self, key: str, ttl: int) -> ThrottlerStorageRecord:
+        """
+        Increment the amount of requests for a given record. The record will
+        automatically be removed from the storage once its TTL has been reached.
+        :param key:
+        :param ttl:
+        :return:
+        """
diff --git a/ellar_throttler/module.py b/ellar_throttler/module.py
@@ -0,0 +1,15 @@
+from ellar.common import Module
+from ellar.core import ModuleBase
+from ellar.di import ProviderConfig
+
+from ellar_throttler.interfaces import IThrottlerStorage
+from ellar_throttler.throttler_service import CacheThrottlerStorageService
+
+
+@Module(
+    providers=[
+        ProviderConfig(IThrottlerStorage, use_class=CacheThrottlerStorageService)
+    ]
+)
+class ThrottlerModule(ModuleBase):
+    pass
diff --git a/ellar_throttler/throttler_guard.py b/ellar_throttler/throttler_guard.py
@@ -0,0 +1,82 @@
+import hashlib
+import typing as t
+
+from ellar.core import GuardCanActivate, IExecutionContext, Response
+from ellar.core.connection import HTTPConnection
+from ellar.di import injectable
+from ellar.helper import get_name
+from ellar.services import Reflector
+
+from .constants import THROTTLER_LIMIT, THROTTLER_SKIP, THROTTLER_TTL
+from .exception import ThrottledException
+from .interfaces import IThrottlerStorage
+
+
+@injectable()
+class ThrottlerGuard(GuardCanActivate):
+    header_prefix = "X-RateLimit"
+
+    def __init__(
+        self, storage_service: IThrottlerStorage, reflector: Reflector
+    ) -> None:
+        self.storage_service = storage_service
+        self.reflector = reflector
+
+    async def can_activate(self, context: IExecutionContext) -> bool:
+        handler = context.get_handler()
+        class_ref = context.get_class()
+
+        # Return early if the current route should be skipped.
+        # or self.options.skipIf?.(context)
+        if self.reflector.get_all_and_override(THROTTLER_SKIP, handler, class_ref):
+            return True
+
+        # Return early when we have no limit or ttl data.
+        route_or_class_limit = self.reflector.get_all_and_override(
+            THROTTLER_LIMIT, handler, class_ref
+        )
+        route_or_class_ttl = self.reflector.get_all_and_override(
+            THROTTLER_TTL, handler, class_ref
+        )
+
+        # Check if specific limits are set at class or route level, otherwise use global options.
+        limit = route_or_class_limit or 60  # or this.options.limit
+        ttl = route_or_class_ttl or 60  # or this.options.ttl
+        return await self.handle_request(context, limit, ttl)
+
+    @classmethod
+    def get_request_response(
+        cls, context: IExecutionContext
+    ) -> t.Tuple[HTTPConnection, Response]:
+        connection_host = context.switch_to_http_connection()
+        return connection_host.get_client(), connection_host.get_response()
+
+    def get_tracker(self, connection: HTTPConnection) -> str:
+        assert connection.client
+        return connection.client.host
+
+    def generate_key(self, context: IExecutionContext, suffix: str) -> str:
+        prefix = f"{get_name(context.get_class())}-{get_name(context.get_handler())}"
+        return hashlib.md5(f"{prefix}-{suffix}".encode("utf8")).hexdigest()
+
+    async def handle_request(
+        self, context: IExecutionContext, limit: int, ttl: int
+    ) -> bool:
+        connection, response = self.get_request_response(context)
+        # TODO: Return early if the current user agent should be ignored.
+
+        tracker = self.get_tracker(connection)
+        key = self.generate_key(context, tracker)
+        result = await self.storage_service.increment(key, ttl)
+
+        # Throw an error when the user reached their limit.
+        if result.total_hits > limit:
+            raise ThrottledException(wait=result.time_to_expire)
+
+        response.headers[f"{self.header_prefix}-Limit"] = str(limit)
+        response.headers[f"{self.header_prefix}-Remaining"] = str(
+            max(0, limit - result.total_hits)
+        )
+        response.headers[f"{self.header_prefix}-Reset"] = str(result.time_to_expire)
+
+        return True
diff --git a/ellar_throttler/throttler_service.py b/ellar_throttler/throttler_service.py
@@ -0,0 +1,103 @@
+import math
+import time
+import typing as t
+
+from ellar.cache import ICacheService
+from ellar.di import injectable
+
+from .interfaces import IThrottlerStorage
+from .throttler_storage_options import ThrottlerStorageOption
+from .throttler_storage_record import ThrottlerStorageRecord
+
+
+@injectable()
+class ThrottlerStorageService(IThrottlerStorage):
+    def __init__(self) -> None:
+        self._storage: t.Dict[str, ThrottlerStorageOption] = {}
+
+    def _set(self, key: str, value: ThrottlerStorageOption) -> None:
+        self._storage[key] = value
+
+    def _has_expired(self, key: str) -> bool:
+        exp = self._storage.get(key)
+        return exp is not None and exp.expires_at <= time.time()
+
+    def _has_key(self, key: str) -> bool:
+        return self._get(key) is not None
+
+    def _delete(self, key: str) -> bool:
+        try:
+            self._storage.pop(key)
+        except KeyError:
+            return False
+        return True
+
+    def _get(self, key: str) -> t.Optional[ThrottlerStorageOption]:
+        if self._has_expired(key):
+            self._delete(key)
+            return None
+
+        return self._storage.get(key)
+
+    @property
+    def storage(self) -> t.Dict[str, ThrottlerStorageOption]:
+        return self._storage
+
+    def get_expiration_time(self, key: str) -> int:
+        return math.floor(self.storage[key].expires_at - time.time())
+
+    async def increment(self, key: str, ttl: int) -> ThrottlerStorageRecord:
+        if not self._has_key(key):
+            self._set(
+                key, ThrottlerStorageOption(total_hits=0, expires_at=time.time() + ttl)
+            )
+
+        history = self._get(key)
+        assert history, "value can not be None"
+        now = time.time()
+
+        time_to_expire = self.get_expiration_time(key)
+
+        if time_to_expire <= 0:
+            history.total_hits = 0
+            history.expires_at = now + ttl
+            time_to_expire = self.get_expiration_time(key)
+
+        history.total_hits += 1
+
+        return ThrottlerStorageRecord(
+            total_hits=history.total_hits, time_to_expire=time_to_expire
+        )
+
+
+@injectable()
+class CacheThrottlerStorageService(IThrottlerStorage):
+    def __init__(self, cache_service: ICacheService) -> None:
+        self._cache_service = cache_service
+
+    @property
+    def storage(self) -> t.Any:
+        return self._cache_service.get_backend()
+
+    async def get_expiration_time(self, key: str) -> int:
+        result = await self._cache_service.get_async(f"{key}-ttl") or 60
+        assert result, "Value can not be none"
+        return math.floor(result - time.time())
+
+    async def increment(self, key: str, ttl: int) -> ThrottlerStorageRecord:
+        if not await self._cache_service.has_key_async(key):
+            await self._cache_service.set_async(key, 0, ttl)
+            await self._cache_service.set_async(f"{key}-ttl", time.time() + ttl, ttl)
+
+        time_to_expire = await self.get_expiration_time(key)
+
+        if time_to_expire <= 0:
+            await self._cache_service.set_async(key, 0, ttl)
+            await self._cache_service.set_async(f"{key}-ttl", time.time() + ttl, ttl)
+            time_to_expire = await self.get_expiration_time(key)
+
+        total_hits = await self._cache_service.incr_async(key)
+
+        return ThrottlerStorageRecord(
+            total_hits=total_hits, time_to_expire=time_to_expire
+        )
diff --git a/ellar_throttler/throttler_storage_options.py b/ellar_throttler/throttler_storage_options.py
@@ -0,0 +1,7 @@
+import dataclasses
+
+
+@dataclasses.dataclass
+class ThrottlerStorageOption:
+    total_hits: int
+    expires_at: float
diff --git a/ellar_throttler/throttler_storage_record.py b/ellar_throttler/throttler_storage_record.py
@@ -0,0 +1,9 @@
+import dataclasses
+
+
+@dataclasses.dataclass
+class ThrottlerStorageRecord:
+    #  Amount of requests done by a specific user (partially based on IP).
+    total_hits: int
+    #  Amount of seconds when the `total_hits` should expire.
+    time_to_expire: int
