Add an initial working in memory rate limiter

This utilises the Generic Cell Rate Algorithm, GCRA, to provide a
(hopefully) elegant rate limiting system.

.gitignore

+*~
+venv/
+__pycache__/
+Quart_Rate_Limiter.egg-info/
+.cache/
+.tox/
+TODO
+.mypy_cache/
+.hypothesis/
+docs/_build/
+.coverage
+.pytest_cache/

.gitlab-ci.yml

+py37:
+  image: python:3.7
+  script:
+  - pip install tox
+  - tox -e format,mypy,py37,pep8,setuppy,manifest

CHANGELOG.rst

+0.1.0
+-----
+
+* Released initial alpha version.

LICENSE

+Copyright P G Jones 2019.
+
+Permission is hereby granted, free of charge, to any person
+obtaining a copy of this software and associated documentation
+files (the "Software"), to deal in the Software without
+restriction, including without limitation the rights to use,
+copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the
+Software is furnished to do so, subject to the following
+conditions:
+
+The above copyright notice and this permission notice shall be
+included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
+OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
+HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
+WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
+OTHER DEALINGS IN THE SOFTWARE.

MANIFEST.in

+include CHANGELOG.rst
+include LICENSE
+include pyproject.toml
+include quart_rate_limiter/py.typed
+include README.rst
+include setup.cfg
+recursive-include quart_limiter *.py
+recursive-include tests *.py
+exclude .gitlab-ci.yml

README.rst

+Quart-Rate-Limiter
+==================
+
+|Build Status| |pypi| |python| |license|
+
+Quart-Rate-Limiter is an extension for `Quart
+<https://gitlab.com/pgjones/quart>`_ to allow for rate limits to be
+defined and enforced on a per route basis.
+
+Usage
+-----
+
+To add a rate limit first initialise the RateLimiting extension with
+the application,
+
+.. code-block:: python
+
+    app = Quart(__name__)
+    rate_limiter = RateLimiter(app)
+
+or via the factory pattern,
+
+.. code-block:: python
+
+    rate_limiter = RateLimiter()
+
+    def create_app():
+        app = Quart(__name__)
+        rate_limiter.init_app(app)
+        return app
+
+Now this is done you can apply rate limits to any route by using the
+``rate_limit`` decorator,
+
+.. code-block:: python
+
+    @app.route('/')
+    @rate_limit(1, timedelta(seconds=10))
+    async def handler():
+        ...
+
+To alter the identification of remote users you can either supply a
+global key function when initialising the extension, or on a per route
+basis.
+
+Simple examples
+~~~~~~~~~~~~~~~
+
+To limit a route to 1 request per second and a maximum of 20 per minute,
+
+.. code-block:: python
+
+    @app.route('/')
+    @rate_limit(1, timedelta(seconds=1))
+    @rate_limit(20, timedelta(minutes=1))
+    async def handler():
+        ...
+
+To identify remote users based on the forwarded IP, rather than the
+direct IP (if behind a load balancer),
+
+.. code-block:: python
+
+    async def key_function():
+        # Return the X-Forwarded-For as the user-agent identifier,
+        # unless it isn't present (direct connection).
+        return request.headers.get("X-Forwarded-For", request.remote_addr)
+
+    RateLimiter(app, key_function=key_function)
+
+The ``key_function`` is a coroutine function to allow session lookups
+if appropriate.
+
+Contributing
+------------
+
+Quart-Rate-Limiter is developed on `GitLab
+<https://gitlab.com/pgjones/quart-rate-limiter>`_. You are very welcome to
+open `issues <https://gitlab.com/pgjones/quart-rate-limiter/issues>`_ or
+propose `merge requests
+<https://gitlab.com/pgjones/quart-rate-limiter/merge_requests>`_.
+
+Testing
+~~~~~~~
+
+The best way to test Quart-Rate-Limiter is with Tox,
+
+.. code-block:: console
+
+    $ pip install tox
+    $ tox
+
+this will check the code style and run the tests.
+
+Help
+----
+
+This README is the best place to start, after that try opening an
+`issue <https://gitlab.com/pgjones/quart-rate-limiter/issues>`_.
+
+
+.. |Build Status| image:: https://gitlab.com/pgjones/quart-rate-limiter/badges/master/build.svg
+   :target: https://gitlab.com/pgjones/quart-rate-limiter/commits/master
+
+.. |pypi| image:: https://img.shields.io/pypi/v/quart-rate-limiter.svg
+   :target: https://pypi.python.org/pypi/Quart-Rate-Limiter/
+
+.. |python| image:: https://img.shields.io/pypi/pyversions/quart-rate-limiter.svg
+   :target: https://pypi.python.org/pypi/Quart-Rate-Limiter/
+
+.. |license| image:: https://img.shields.io/badge/license-MIT-blue.svg
+   :target: https://gitlab.com/pgjones/quart-rate-limiter/blob/master/LICENSE

pyproject.toml

+[build-system]
+requires = ["setuptools", "wheel"]
+
+[tool.black]
+line-length = 100
+target-version = ["py37"]
+
+[tool.isort]
+dont_skip = ["__init__.py"]
+combine_as_imports = true
+force_grid_wrap = 0
+include_trailing_comma = true
+known_first_party = "quart_rate_limiter, tests"
+known_third_party = "_pytest, pytest, quart"
+line_length = 100
+multi_line_output = 3
+no_lines_before = "LOCALFOLDER"
+order_by_type = false
+reverse_relative = true

quart_rate_limiter/__init__.py

+from abc import ABCMeta, abstractmethod
+from dataclasses import dataclass
+from datetime import datetime, timedelta
+from typing import Awaitable, Callable, Dict, List, Optional, Type
+
+from quart import current_app, Quart, request, Response
+from quart.exceptions import HTTPException
+
+QUART_RATE_LIMITER_ATTRIBUTE = "_quart_rate_limiter_limits"
+
+KeyCallable = Callable[[], Awaitable[str]]
+
+
+class RateLimitExceeded(HTTPException):
+    """A 429 Rate limit exceeded error.
+
+    Arguments:
+        retry_after: Seconds left till the remaining resets to the limit.
+    """
+
+    def __init__(self, retry_after: int) -> None:
+        super().__init__(429, "Rate Limit Exceeded", "RATE_LIMIT_EXCEEDED")
+        self.retry_after = retry_after
+
+    def get_headers(self) -> dict:
+        return {"Retry-After": str(self.retry_after)}
+
+
+class RateLimiterStoreABC(metaclass=ABCMeta):
+    @abstractmethod
+    async def get(self, key: str, default: datetime) -> datetime:
+        """Get the TAT for the given *key* if present or the *default* if not.
+
+        Arguments:
+            key: The key to indentify the TAT.
+            default: If no TAT for the *key* is available, return this
+                default.
+
+        Returns:
+            A Theoretical Arrival Time, TAT.
+        """
+        pass
+
+    @abstractmethod
+    async def set(self, key: str, tat: datetime) -> None:
+        """Set the TAT for the given *key*.
+
+        Arguments:
+            key: The key to indentify the TAT.
+            tat: The TAT value to set.
+        """
+        pass
+
+
+class MemoryStore(RateLimiterStoreABC):
+    """An in memory store of rate limits."""
+
+    def __init__(self) -> None:
+        self._tats: Dict[str, datetime] = {}
+
+    async def get(self, key: str, default: datetime) -> datetime:
+        return self._tats.get(key, default)
+
+    async def set(self, key: str, tat: datetime) -> None:
+        self._tats[key] = tat
+
+
+@dataclass
+class RateLimit:
+    count: int
+    period: timedelta
+    key_function: Optional[KeyCallable]
+
+    @property
+    def inverse(self) -> float:
+        return self.period.total_seconds() / self.count
+
+    @property
+    def key(self) -> str:
+        return f"{self.count}-{self.period.total_seconds()}"
+
+
+def rate_limit(
+    limit: int, period: timedelta, key_function: Optional[KeyCallable] = None
+) -> Callable:
+    """A decorator to add a rate limit marker to the route.
+
+    This should be used to wrap a route handler (or view function) to
+    apply a rate limit to requests to that route. Note that it is
+    important that this decorator be wrapped by the route decorator
+    and not vice, versa, as below.
+
+    .. code-block:: python
+
+        @app.route('/')
+        @rate_limit(10, timedelta(seconds=10))
+        async def index():
+            ...
+
+    Arguments:
+        limit: The maximum number of requests to serve within a
+            period.
+        period: The duration over which the number of requests must
+            not exceed the *limit*.
+        key_function: A coroutine function that returns a unique key
+            to identify the user agent.
+
+    .. code-block:: python
+
+        async def example_key_function() -> str:
+            return request.remote_addr
+
+    """
+
+    def decorator(func: Callable) -> Callable:
+        rate_limits = getattr(func, QUART_RATE_LIMITER_ATTRIBUTE, [])
+        rate_limits.append(RateLimit(limit, period, key_function))
+        setattr(func, QUART_RATE_LIMITER_ATTRIBUTE, rate_limits)
+        return func
+
+    return decorator
+
+
+async def remote_addr_key() -> str:
+    return request.remote_addr
+
+
+class RateLimiter:
+    """A Rate limiter instance.
+
+    This can be used to initialise Rate Limiting for a given app,
+    either directly,
+
+    .. code-block:: python
+
+        app = Quart(__name__)
+        rate_limiter = RateLimiter(app)
+
+    or via the factory pattern,
+
+    .. code-block:: python
+
+        rate_limiter = RateLimiter()
+
+        def create_app():
+            app = Quart(__name__)
+            rate_limiter.init_app(app)
+            return app
+
+    The limiter itself can be customised using the following
+    arguments,
+
+    Arguments:
+        key_function: A coroutine function that returns a unique key
+            to identify the user agent.
+        rate_limit_exception: A type of exception to raise if the rate
+            limit has been exceeded.
+        store: The store that contains the theoretical arrival times by
+            key.
+    """
+
+    def __init__(
+        self,
+        app: Optional[Quart] = None,
+        key_function: KeyCallable = remote_addr_key,
+        rate_limit_exception: Type[Exception] = RateLimitExceeded,
+        store: Optional[RateLimiterStoreABC] = None,
+    ) -> None:
+        self.key_function = key_function
+        self.rate_limit_exception = rate_limit_exception
+        self.store: RateLimiterStoreABC
+        if store is None:
+            self.store = MemoryStore()
+        else:
+            self.store = store
+
+        if app is not None:
+            self.init_app(app)
+
+    def init_app(self, app: Quart) -> None:
+        app.before_request(self._before_request)
+        app.after_request(self._after_request)
+
+    async def _before_request(self) -> None:
+        endpoint = request.endpoint
+        view_func = current_app.view_functions.get(endpoint)
+        if view_func is not None:
+            rate_limits: List[RateLimit] = getattr(view_func, QUART_RATE_LIMITER_ATTRIBUTE, [])
+            await self._raise_on_rejection(endpoint, rate_limits)
+            await self._update_limits(endpoint, rate_limits)
+
+    async def _raise_on_rejection(self, endpoint: str, rate_limits: List[RateLimit]) -> None:
+        now = datetime.utcnow()
+        for rate_limit in rate_limits:
+            key = await self._create_key(endpoint, rate_limit)
+            # This is the GCRA rate limiting system and tat stands for
+            # the theoretical arrival time.
+            tat = max(await self.store.get(key, now), now)
+            separation = (tat - now).total_seconds()
+            max_interval = rate_limit.period.total_seconds() - rate_limit.inverse
+            if separation > max_interval:
+                retry_after = ((tat - timedelta(seconds=max_interval)) - now).total_seconds()
+                raise self.rate_limit_exception(int(retry_after))
+
+    async def _update_limits(self, endpoint: str, rate_limits: List[RateLimit]) -> None:
+        # Update the tats for all the rate limits. This must only
+        # occur if no limit rejects the request.
+        now = datetime.utcnow()
+        for rate_limit in rate_limits:
+            key = await self._create_key(endpoint, rate_limit)
+            tat = max(await self.store.get(key, now), now)
+            new_tat = max(tat, now) + timedelta(seconds=rate_limit.inverse)
+            await self.store.set(key, new_tat)
+
+    async def _after_request(self, response: Response) -> Response:
+        endpoint = request.endpoint
+        view_func = current_app.view_functions.get(endpoint)
+        rate_limits: List[RateLimit] = getattr(view_func, QUART_RATE_LIMITER_ATTRIBUTE, [])
+        try:
+            min_limit = min(rate_limits, key=lambda rate_limit: rate_limit.period.total_seconds())
+        except ValueError:
+            pass  # No rate limits
+        else:
+            key = await self._create_key(endpoint, min_limit)
+            now = datetime.utcnow()
+            tat = max(await self.store.get(key, now), now)
+            separation = (tat - now).total_seconds()
+            remaining = int((min_limit.period.total_seconds() - separation) / min_limit.inverse)
+            max_interval = min_limit.period.total_seconds() - min_limit.inverse
+            reset = int(((tat - timedelta(seconds=max_interval)) - now).total_seconds())
+            response.headers["RateLimit-Limit"] = str(min_limit.count)
+            response.headers["RateLimit-Remaining"] = str(remaining)
+            response.headers["RateLimit-Reset"] = str(reset)
+        return response
+
+    async def _create_key(self, endpoint: str, rate_limit: RateLimit) -> str:
+        key_function = rate_limit.key_function or self.key_function
+        key = await key_function()
+        app_name = current_app.import_name
+        return f"{app_name}-{endpoint}-{rate_limit.key}-{key}"

quart_rate_limiter/py.typed

+Marker

setup.cfg

+[check-manifest]
+ignore = tox.ini
+
+[flake8]
+ignore = E203, E252, W503, W504
+max_line_length = 100
+
+[mypy]
+allow_redefinition = True
+disallow_subclassing_any = True
+disallow_untyped_defs = True
+strict_equality = True
+strict_optional = False
+warn_redundant_casts = True
+warn_unused_configs = True
+warn_unused_ignores = True
+
+[mypy-_pytest.*]
+ignore_missing_imports = True
+
+[mypy-pytest.*]
+ignore_missing_imports = True
+
+[tool:pytest]
+addopts = --showlocals

setup.py

+import os
+import sys
+
+from setuptools import setup, find_packages
+
+if sys.version_info < (3,7,0):
+    sys.exit("Python 3.7.0 is the minimum required version")
+
+PROJECT_ROOT = os.path.dirname(__file__)
+
+with open(os.path.join(PROJECT_ROOT, "README.rst")) as file_:
+    long_description = file_.read()
+
+INSTALL_REQUIRES = [
+    "Quart>=0.7",
+]
+
+setup(
+    name="Quart-Rate-Limiter",
+    version="0.1.0",
+    python_requires=">=3.7.0",
+    description="A Quart extension to provide rate limiting support.",
+    long_description=long_description,
+    url="https://gitlab.com/pgjones/quart-rate-limiter/",
+    author="P G Jones",
+    author_email="philip.graham.jones@googlemail.com",
+    license="MIT",
+    classifiers=[
+        "Development Status :: 3 - Alpha",
+        "Environment :: Web Environment",
+        "Intended Audience :: Developers",
+        "License :: OSI Approved :: MIT License",
+        "Operating System :: OS Independent",
+        "Programming Language :: Python",
+        "Programming Language :: Python :: 3",
+        "Programming Language :: Python :: 3.7",
+        "Topic :: Internet :: WWW/HTTP :: Dynamic Content",
+        "Topic :: Software Development :: Libraries :: Python Modules",
+    ],
+    packages=find_packages(exclude=["tests", "tests.*"]),
+    py_modules=["quart_rate_limiter"],
+    install_requires=INSTALL_REQUIRES,
+    tests_require=INSTALL_REQUIRES + [
+        "pytest",
+        "pytest-asyncio",
+    ],
+    include_package_data=True,
+)

tests/test_basic.py

+from datetime import datetime, timedelta
+
+import pytest
+from _pytest.monkeypatch import MonkeyPatch
+from quart import Quart, ResponseReturnValue
+
+import quart_rate_limiter
+from quart_rate_limiter import rate_limit, RateLimiter
+
+
+@pytest.fixture(name="fixed_datetime")
+def _fixed_datetime(monkeypatch: MonkeyPatch) -> datetime:
+    class MockDatetime(datetime):
+        @classmethod
+        def utcnow(cls) -> datetime:
+            return datetime(2019, 3, 4)
+
+    monkeypatch.setattr(quart_rate_limiter, "datetime", MockDatetime)
+    return MockDatetime.utcnow()
+
+
+@pytest.fixture(name="app")
+def _app() -> Quart:
+    app = Quart(__name__)
+
+    @app.route("/rate_limit/")
+    @rate_limit(1, timedelta(seconds=2))
+    @rate_limit(10, timedelta(seconds=20))
+    async def index() -> ResponseReturnValue:
+        return ""
+
+    RateLimiter(app)
+    return app
+
+
+@pytest.mark.asyncio
+async def test_rate_limit(app: Quart, fixed_datetime: datetime) -> None:
+    test_client = app.test_client()
+    response = await test_client.get("/rate_limit/")
+    assert response.status_code == 200
+    assert response.headers["RateLimit-Limit"] == "1"
+    assert response.headers["RateLimit-Remaining"] == "0"
+    assert response.headers["RateLimit-Reset"] == "2"
+
+    response = await test_client.get("/rate_limit/")
+    assert response.status_code == 429
+    assert response.headers["Retry-After"] == "2"
+
+
+@pytest.mark.asyncio
+async def test_rate_limit_unique_keys(app: Quart, fixed_datetime: datetime) -> None:
+    test_client = app.test_client()
+    response = await test_client.get("/rate_limit/", headers={"Remote-Addr": "127.0.0.1"})
+    assert response.status_code == 200
+    response = await test_client.get("/rate_limit/", headers={"Remote-Addr": "127.0.0.2"})
+    assert response.status_code == 200

tox.ini

+[tox]
+envlist = format,mypy,py37,pep8,setuppy,manifest
+
+[testenv]
+deps =
+    pytest
+    pytest-asyncio
+    pytest-cov
+    pytest-sugar
+commands = pytest --cov=quart_rate_limiter tests/
+
+[testenv:format]
+basepython = python3.7
+deps =
+    black
+    isort
+commands =
+    black --check --diff quart_rate_limiter/ tests/
+    isort --check --diff --recursive quart_rate_limiter/ tests
+
+[testenv:pep8]
+basepython = python3.7
+deps =
+    flake8
+    pep8-naming
+    flake8-print
+commands = flake8 quart_rate_limiter/ tests/
+
+[testenv:mypy]
+basepython = python3.7
+deps = mypy
+commands =
+    mypy quart_rate_limiter/ tests/
+
+[testenv:setuppy]
+basepython = python3.7
+deps =
+    docutils
+    Pygments
+commands =
+    python setup.py check \
+        --metadata \
+        --restructuredtext \
+        --strict
+
+[testenv:manifest]
+basepython = python3.7
+deps = check-manifest
+commands = check-manifest