Rate limiting for scheduling triggers

flexmeasures/api/v3_0/__init__.py

@@ -11,7 +11,10 @@ def register_at(app: Flask):
 
     v3_0_api_prefix = "/api/v3_0"
 
-    SensorAPI.register(app, route_prefix=v3_0_api_prefix)
-    UserAPI.register(app, route_prefix=v3_0_api_prefix)
-    AssetAPI.register(app, route_prefix=v3_0_api_prefix)
-    HealthAPI.register(app, route_prefix=v3_0_api_prefix)
+    apis = (SensorAPI, UserAPI, AssetAPI, HealthAPI)
+
+    for api_cls in apis:
+        if hasattr(api_cls, "_rate_limiter"):
+            # function will return a flask-rate limit defined by this API
+            api_cls.decorators.append(api_cls._rate_limiter(app))
+        api_cls.register(app, route_prefix=v3_0_api_prefix)

flexmeasures/api/v3_0/sensors.py

@@ -1,10 +1,10 @@
 from datetime import datetime, timedelta
 from typing import Optional
 
-from flask import current_app
+from flask import current_app, Flask, request
 from flask_classful import FlaskView, route
 from flask_json import as_json
-from flask_security import auth_required
+from flask_security import auth_required, current_user
 import isodate
 from marshmallow import fields, ValidationError
 from rq.job import Job, NoSuchJobError
@@ -53,6 +53,15 @@ class SensorAPI(FlaskView):
     trailing_slash = False
     decorators = [auth_required()]
 
+    @classmethod
+    def _rate_limiter(cls, app: Flask):
+        # a schedule can be triggered once per 5 minutes per account & sensor
+        return app.limiter.limit(
+            "1 per 5 minutes",
+            key_func=lambda: f"{current_user.account_id}-{request.view_args.get('id', '')}",
+            cost=lambda: 1 if request.endpoint == "SensorAPI:trigger_schedule" else 0,
+        )
+
     @route("", methods=["GET"])
     @use_kwargs(
         {

flexmeasures/api/v3_0/tests/test_sensor_schedules.py

@@ -113,6 +113,15 @@ def test_trigger_and_get_schedule(
         assert trigger_schedule_response.status_code == 200
         job_id = trigger_schedule_response.json["schedule"]
 
+        # Check rate limiter kicks in when immediately reattempting to trigger a schedule
+        trigger_schedule_response = client.post(
+            url_for("SensorAPI:trigger_schedule", id=sensor.id),
+            json=message,
+            headers={"Authorization": auth_token},
+        )
+        print("Server responded with:\n%s" % trigger_schedule_response.json)
+        assert trigger_schedule_response.status_code == 429
+
     # look for scheduling jobs in queue
     assert (
         len(app.queues["scheduling"]) == 1
@@ -205,6 +214,16 @@ def test_trigger_and_get_schedule(
     )
     print("Server responded with:\n%s" % get_schedule_response.json)
     assert get_schedule_response.status_code == 200
+
+    # Check rate limiter does not kick in when immediately reattempting to fetch a schedule
+    get_schedule_response = client.get(
+        url_for("SensorAPI:get_schedule", id=sensor.id, uuid=job_id),
+        query_string=get_schedule_message,
+        headers={"content-type": "application/json", "Authorization": auth_token},
+    )
+    print("Server responded with:\n%s" % get_schedule_response.json)
+    assert get_schedule_response.status_code == 200  # not 429
+
     # assert get_schedule_response.json["type"] == "GetDeviceMessageResponse"
     assert len(get_schedule_response.json["values"]) == expected_length_of_schedule
 

flexmeasures/app.py

@@ -5,6 +5,8 @@
 
 from flask import Flask, g, request
 from flask.cli import load_dotenv
+from flask_limiter import Limiter
+from flask_limiter.util import get_remote_address
 from flask_mail import Mail
 from flask_sslify import SSLify
 from flask_json import FlaskJSON
@@ -67,10 +69,7 @@ def create(
     if app.testing:
         from fakeredis import FakeStrictRedis
 
-        app.queues = dict(
-            forecasting=Queue(connection=FakeStrictRedis(), name="forecasting"),
-            scheduling=Queue(connection=FakeStrictRedis(), name="scheduling"),
-        )
+        redis_conn = FakeStrictRedis()
     else:
         redis_conn = Redis(
             app.config["FLEXMEASURES_REDIS_URL"],
@@ -83,10 +82,20 @@ def create(
             redis_conn = Redis("MY-DB-NAME", unix_socket_path="/tmp/my-redis.socket",
             )
         """
-        app.queues = dict(
-            forecasting=Queue(connection=redis_conn, name="forecasting"),
-            scheduling=Queue(connection=redis_conn, name="scheduling"),
-        )
+    app.queues = dict(
+        forecasting=Queue(connection=redis_conn, name="forecasting"),
+        scheduling=Queue(connection=redis_conn, name="scheduling"),
+    )
+
+    # Set up rate limiter
+    app.config["RATELIMIT_STORAGE_URI"] = "redis://"
+    app.config["RATELIMIT_STORAGE_OPTIONS"] = {
+        "connection_pool": redis_conn.connection_pool
+    }
+    app.limiter = Limiter(
+        get_remote_address,
+        app=app,
+    )
 
     # Some basic security measures
 

requirements/app.in

@@ -37,6 +37,7 @@ importlib_metadata
 sqlalchemy>=1.4.0, <2
 Flask-SSLify
 Flask_JSON
+Flask-Limiter[redis]
 Flask-Migrate
 Flask-WTF
 Flask-Mail

requirements/app.txt

@@ -40,14 +40,19 @@ click-default-group==1.2.2
     # via -r requirements/app.in
 colour==0.1.5
     # via -r requirements/app.in
+commonmark==0.9.1
+    # via rich
 contourpy==1.0.7
     # via matplotlib
 convertdate==2.4.0
     # via workalendar
 cycler==0.11.0
     # via matplotlib
 deprecated==1.2.13
-    # via sktime
+    # via
+    #   limits
+    #   redis
+    #   sktime
 dill==0.3.6
     # via openturns
 dnspython==2.3.0
@@ -66,6 +71,7 @@ flask==2.1.2
     #   flask-classful
     #   flask-cors
     #   flask-json
+    #   flask-limiter
     #   flask-login
     #   flask-mail
     #   flask-marshmallow
@@ -83,6 +89,8 @@ flask-cors==3.0.10
     # via -r requirements/app.in
 flask-json==0.3.5
     # via -r requirements/app.in
+flask-limiter[redis]==3.1.0
+    # via -r requirements/app.in
 flask-login==0.6.1
     # via
     #   -r requirements/app.in
@@ -147,6 +155,8 @@ jsonschema==4.17.3
     # via altair
 kiwisolver==1.4.4
     # via matplotlib
+limits==3.1.6
+    # via flask-limiter
 llvmlite==0.39.1
     # via numba
 lunardate==0.2.0
@@ -192,8 +202,11 @@ numpy==1.23.5
     #   uniplot
 openturns==1.20.post3
     # via timely-beliefs
+ordered-set==4.1.0
+    # via flask-limiter
 packaging==23.0
     # via
+    #   limits
     #   marshmallow
     #   marshmallow-sqlalchemy
     #   matplotlib
@@ -231,6 +244,8 @@ py-moneyed==3.0
     # via -r requirements/app.in
 pydantic==1.10.6
     # via inflect
+pygments==2.14.0
+    # via rich
 pyluach==2.2.0
     # via workalendar
 pymeeus==0.5.12
@@ -267,6 +282,8 @@ requests==2.28.2
     #   tldextract
 requests-file==1.5.1
     # via tldextract
+rich==12.6.0
+    # via flask-limiter
 rq==1.13.0
     # via
     #   -r requirements/app.in
@@ -321,6 +338,8 @@ toolz==0.12.0
     # via altair
 typing-extensions==4.5.0
     # via
+    #   flask-limiter
+    #   limits
     #   alembic
     #   py-moneyed
     #   pydantic

requirements/docs.txt

@@ -43,6 +43,7 @@ packaging==23.0
     #   sphinx
 pygments==2.14.0
     # via
+    #   -c requirements/app.txt
     #   sphinx
     #   sphinx-tabs
 requests==2.28.2