Refactor rate limiting (fixes #326) (#327)

DavidMStraub · web-flow · commit e8e6d1ef86fa · 2022-12-29T17:08:24.000+01:00
* Refactor rate limiting

* Unit test for rate limiter
diff --git a/gramps_webapi/api/ratelimiter.py b/gramps_webapi/api/ratelimiter.py
@@ -0,0 +1,7 @@
+"""Rate limiting decorator."""
+
+
+from flask_limiter import Limiter
+from flask_limiter.util import get_remote_address
+
+limiter = Limiter(key_func=get_remote_address)
diff --git a/gramps_webapi/api/resources/config.py b/gramps_webapi/api/resources/config.py
@@ -19,22 +19,15 @@
 
 """User administration resources."""
 
-import datetime
-from gettext import gettext as _
 
-from flask import abort, current_app, jsonify, render_template
-from flask_jwt_extended import create_access_token, get_jwt, get_jwt_identity
-from flask_limiter import Limiter
-from flask_limiter.util import get_remote_address
+from flask import abort, current_app, jsonify
 from webargs import fields
 
 from ...auth.const import PERM_EDIT_SETTINGS, PERM_VIEW_SETTINGS
 from ...const import DB_CONFIG_ALLOWED_KEYS
 from ..auth import require_permissions
 from ..util import use_args
-from . import ProtectedResource, Resource
-
-limiter = Limiter(key_func=get_remote_address)
+from . import ProtectedResource
 
 
 class ConfigsResource(ProtectedResource):
diff --git a/gramps_webapi/api/resources/token.py b/gramps_webapi/api/resources/token.py
@@ -19,7 +19,6 @@
 
 """Authentication endpoint blueprint."""
 
-import datetime
 from typing import Iterable
 
 from flask import abort, current_app
@@ -28,16 +27,13 @@
     create_refresh_token,
     get_jwt_identity,
 )
-from flask_limiter import Limiter
-from flask_limiter.util import get_remote_address
 from webargs import fields, validate
 
 from ...auth.const import CLAIM_LIMITED_SCOPE, SCOPE_CREATE_OWNER
+from ..ratelimiter import limiter
 from ..util import use_args
 from . import RefreshProtectedResource, Resource
 
-limiter = Limiter(key_func=get_remote_address)
-
 
 def get_tokens(user_id: str, permissions: Iterable[str], include_refresh: bool = False):
     """Create access token (and refresh token if desired)."""
diff --git a/gramps_webapi/api/resources/user.py b/gramps_webapi/api/resources/user.py
@@ -24,8 +24,6 @@
 
 from flask import abort, current_app, jsonify, render_template
 from flask_jwt_extended import create_access_token, get_jwt, get_jwt_identity
-from flask_limiter import Limiter
-from flask_limiter.util import get_remote_address
 from webargs import fields
 
 from ...auth.const import (
@@ -44,6 +42,7 @@
     SCOPE_RESET_PW,
 )
 from ..auth import require_permissions
+from ..ratelimiter import limiter
 from ..tasks import (
     send_email_confirm_email,
     send_email_new_user,
@@ -52,8 +51,6 @@
 from ..util import use_args
 from . import LimitedScopeProtectedResource, ProtectedResource, Resource
 
-limiter = Limiter(key_func=get_remote_address)
-
 
 class UserChangeBase(ProtectedResource):
     """Base class for user change endpoints."""
diff --git a/gramps_webapi/app.py b/gramps_webapi/app.py
@@ -21,6 +21,7 @@
 
 import logging
 import os
+from typing import Any, Dict, Optional
 
 from flask import Flask, abort, g, send_from_directory
 from flask_compress import Compress
@@ -29,15 +30,17 @@
 
 from .api import api_blueprint
 from .api.cache import thumbnail_cache
-from .api.resources.token import limiter
+from .api.ratelimiter import limiter
 from .api.search import SearchIndexer
 from .auth import SQLAuth
 from .config import DefaultConfig, DefaultConfigJWT
 from .const import API_PREFIX, ENV_CONFIG_FILE
 from .dbmanager import WebDbManager
 
 
-def create_app(db_manager=None):
+def create_app(
+    db_manager: Optional[WebDbManager] = None, config: Optional[Dict[str, Any]] = None
+):
     """Flask application factory."""
     app = Flask(__name__)
 
@@ -97,6 +100,10 @@ def create_app(db_manager=None):
         "MEDIA_BASE_DIR"
     )
 
+    # update config from dictionary if present
+    if config:
+        app.config.update(**config)
+
     # instantiate DB manager
     if db_manager is None:
         app.config["DB_MANAGER"] = WebDbManager(
diff --git a/setup.py b/setup.py
@@ -35,7 +35,7 @@
     "Flask-Compress",
     "Flask-Cors",
     "Flask-JWT-Extended>=4.2.1, !=4.4.0, !=4.4.1",
-    "Flask-Limiter<=2.8.0",
+    "Flask-Limiter>=2.9.0",
     "marshmallow>=3.13.0",
     "webargs",
     "SQLAlchemy",
diff --git a/tests/test_endpoints/__init__.py b/tests/test_endpoints/__init__.py
@@ -67,8 +67,9 @@ def setUpModule():
 
     test_db = ExampleDbSQLite()
     with patch.dict("os.environ", {ENV_CONFIG_FILE: TEST_EXAMPLE_GRAMPS_AUTH_CONFIG}):
-        test_app = create_app(db_manager=test_db)
-    test_app.config["TESTING"] = True
+        test_app = create_app(
+            db_manager=test_db, config={"TESTING": True, "RATELIMIT_ENABLED": False}
+        )
     TEST_CLIENT = test_app.test_client()
     search_index = test_app.config["SEARCH_INDEXER"]
     db = test_db.get_db().db
diff --git a/tests/test_endpoints/test_config.py b/tests/test_endpoints/test_config.py
@@ -44,8 +44,7 @@ def setUp(self):
         self.dbman = CLIDbManager(DbState())
         _, _name = self.dbman.create_new_db_cli(self.name, dbid="sqlite")
         with patch.dict("os.environ", {ENV_CONFIG_FILE: TEST_AUTH_CONFIG}):
-            self.app = create_app()
-        self.app.config["TESTING"] = True
+            self.app = create_app(config={"TESTING": True, "RATELIMIT_ENABLED": False})
         self.client = self.app.test_client()
         sqlauth = self.app.config["AUTH_PROVIDER"]
         sqlauth.create_table()
diff --git a/tests/test_endpoints/test_user.py b/tests/test_endpoints/test_user.py
@@ -46,8 +46,7 @@ def setUp(self):
         self.dbman = CLIDbManager(DbState())
         _, _name = self.dbman.create_new_db_cli(self.name, dbid="sqlite")
         with patch.dict("os.environ", {ENV_CONFIG_FILE: TEST_AUTH_CONFIG}):
-            self.app = create_app()
-        self.app.config["TESTING"] = True
+            self.app = create_app(config={"TESTING": True, "RATELIMIT_ENABLED": False})
         self.client = self.app.test_client()
         sqlauth = self.app.config["AUTH_PROVIDER"]
         sqlauth.create_table()
@@ -703,8 +702,7 @@ def setUp(self):
         self.dbman = CLIDbManager(DbState())
         _, _name = self.dbman.create_new_db_cli(self.name, dbid="sqlite")
         with patch.dict("os.environ", {ENV_CONFIG_FILE: TEST_AUTH_CONFIG}):
-            self.app = create_app()
-        self.app.config["TESTING"] = True
+            self.app = create_app(config={"TESTING": True, "RATELIMIT_ENABLED": False})
         self.client = self.app.test_client()
         sqlauth = self.app.config["AUTH_PROVIDER"]
         sqlauth.create_table()
diff --git a/tests/test_jwt.py b/tests/test_jwt.py
@@ -53,8 +53,7 @@ def setUpClass(cls):
         cls.dbman = CLIDbManager(DbState())
         _, _name = cls.dbman.create_new_db_cli(cls.name, dbid="sqlite")
         with patch.dict("os.environ", {ENV_CONFIG_FILE: TEST_AUTH_CONFIG}):
-            cls.app = create_app()
-        cls.app.config["TESTING"] = True
+            cls.app = create_app(config={"TESTING": True, "RATELIMIT_ENABLED": False})
         cls.client = cls.app.test_client()
         sqlauth = cls.app.config["AUTH_PROVIDER"]
         sqlauth.create_table()
@@ -94,6 +93,7 @@ def test_person_endpoint(self):
         rv = self.client.post(
             "/api/token/", json={"username": "user", "password": "123"}
         )
+        assert rv.status_code == 200
         token = rv.json["access_token"]
         rv = self.client.get(
             "/api/people/" + it["handle"] + "?profile=all",
@@ -114,6 +114,7 @@ def test_person_endpoint_privacy(self):
         rv = self.client.post(
             "/api/token/", json={"username": "user", "password": "123"}
         )
+        assert rv.status_code == 200
         token_user = rv.json["access_token"]
         rv = self.client.post(
             "/api/token/", json={"username": "admin", "password": "123"}
@@ -159,6 +160,7 @@ def test_refresh_token_endpoint(self):
         rv = self.client.post(
             "/api/token/", json={"username": "user", "password": "123"}
         )
+        assert rv.status_code == 200
         refresh_token = rv.json["refresh_token"]
         access_token = rv.json["access_token"]
         # incorrectly send access token instead of refresh token!
diff --git a/tests/test_limiter.py b/tests/test_limiter.py
@@ -0,0 +1,66 @@
+#
+# Gramps Web API - A RESTful API for the Gramps genealogy program
+#
+# Copyright (C) 2022      David Straub
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation; either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program. If not, see <https://www.gnu.org/licenses/>.
+#
+
+"""Tests for rate limited endpoints."""
+
+import unittest
+from unittest.mock import patch
+
+from time import sleep
+
+from gramps.cli.clidbman import CLIDbManager
+from gramps.gen.dbstate import DbState
+
+from gramps_webapi.app import create_app
+from gramps_webapi.auth.const import ROLE_GUEST, ROLE_OWNER
+from gramps_webapi.const import ENV_CONFIG_FILE, TEST_AUTH_CONFIG
+
+
+class TestRateLimits(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.name = "Test Web API"
+        cls.dbman = CLIDbManager(DbState())
+        _, _name = cls.dbman.create_new_db_cli(cls.name, dbid="sqlite")
+        with patch.dict("os.environ", {ENV_CONFIG_FILE: TEST_AUTH_CONFIG}):
+            cls.app = create_app(config={"TESTING": True, "RATELIMIT_ENABLED": True})
+        cls.client = cls.app.test_client()
+        sqlauth = cls.app.config["AUTH_PROVIDER"]
+        sqlauth.create_table()
+        sqlauth.add_user(name="user", password="123", role=ROLE_GUEST)
+        sqlauth.add_user(name="admin", password="123", role=ROLE_OWNER)
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.dbman.remove_database(cls.name)
+
+    def test_token_endpoint(self):
+        rv = self.client.post(
+            "/api/token/", json={"username": "user", "password": "123"}
+        )
+        assert rv.status_code == 200
+        rv = self.client.post(
+            "/api/token/", json={"username": "user", "password": "123"}
+        )
+        assert rv.status_code == 429
+        sleep(1)
+        rv = self.client.post(
+            "/api/token/", json={"username": "user", "password": "123"}
+        )
+        assert rv.status_code == 200
