From ce7f87c3763569f0e5319b2013c66d0daabfa7aa Mon Sep 17 00:00:00 2001 From: "F.N. Claessen" Date: Wed, 18 Jan 2023 22:10:00 +0100 Subject: [PATCH 1/4] Tech spike: rate limiting for scheduling triggers Signed-off-by: F.N. Claessen --- flexmeasures/api/v3_0/__init__.py | 11 ++++++++- .../api/v3_0/tests/test_sensor_schedules.py | 19 +++++++++++++++ flexmeasures/app.py | 23 ++++++++++++------- requirements/app.in | 1 + 4 files changed, 45 insertions(+), 9 deletions(-) diff --git a/flexmeasures/api/v3_0/__init__.py b/flexmeasures/api/v3_0/__init__.py index 56bc4d8401..92bed2044b 100644 --- a/flexmeasures/api/v3_0/__init__.py +++ b/flexmeasures/api/v3_0/__init__.py @@ -1,4 +1,5 @@ -from flask import Flask +from flask import Flask, request +from flask_security import current_user from flexmeasures.api.v3_0.sensors import SensorAPI from flexmeasures.api.v3_0.users import UserAPI @@ -11,6 +12,14 @@ def register_at(app: Flask): v3_0_api_prefix = "/api/v3_0" + def cost_function() -> int: + if request.endpoint == "SensorAPI:trigger_schedule": + return 1 + return 0 + + # Apply rate limit: a schedule can be triggered once per 5 minutes per sensor per account + SensorAPI.decorators.append(app.limiter.limit("1 per 5 minutes", key_func=lambda: str(current_user.account_id) + request.view_args.get("id", ""), cost=cost_function)) + 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) diff --git a/flexmeasures/api/v3_0/tests/test_sensor_schedules.py b/flexmeasures/api/v3_0/tests/test_sensor_schedules.py index 2c42ae5609..0fea39453d 100644 --- a/flexmeasures/api/v3_0/tests/test_sensor_schedules.py +++ b/flexmeasures/api/v3_0/tests/test_sensor_schedules.py @@ -105,6 +105,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 @@ -197,6 +206,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 diff --git a/flexmeasures/app.py b/flexmeasures/app.py index 7f0a450c7f..3b6b2303c1 100644 --- a/flexmeasures/app.py +++ b/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,18 @@ 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 diff --git a/requirements/app.in b/requirements/app.in index 90cd59ef29..73267480bc 100644 --- a/requirements/app.in +++ b/requirements/app.in @@ -36,6 +36,7 @@ importlib_metadata sqlalchemy>=1.4.0 Flask-SSLify Flask_JSON +Flask-Limiter[redis] Flask-SQLAlchemy>=2.4.3 Flask-Migrate Flask-WTF From c393e7595b5457675f040807982bddbe56ae5fe6 Mon Sep 17 00:00:00 2001 From: "F.N. Claessen" Date: Thu, 19 Jan 2023 10:09:15 +0100 Subject: [PATCH 2/4] black Signed-off-by: F.N. Claessen --- flexmeasures/api/v3_0/__init__.py | 9 ++++++++- flexmeasures/app.py | 4 +++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/flexmeasures/api/v3_0/__init__.py b/flexmeasures/api/v3_0/__init__.py index 92bed2044b..2a060bbe2e 100644 --- a/flexmeasures/api/v3_0/__init__.py +++ b/flexmeasures/api/v3_0/__init__.py @@ -18,7 +18,14 @@ def cost_function() -> int: return 0 # Apply rate limit: a schedule can be triggered once per 5 minutes per sensor per account - SensorAPI.decorators.append(app.limiter.limit("1 per 5 minutes", key_func=lambda: str(current_user.account_id) + request.view_args.get("id", ""), cost=cost_function)) + SensorAPI.decorators.append( + app.limiter.limit( + "1 per 5 minutes", + key_func=lambda: str(current_user.account_id) + + request.view_args.get("id", ""), + cost=cost_function, + ) + ) SensorAPI.register(app, route_prefix=v3_0_api_prefix) UserAPI.register(app, route_prefix=v3_0_api_prefix) diff --git a/flexmeasures/app.py b/flexmeasures/app.py index 3b6b2303c1..20fb4d62a5 100644 --- a/flexmeasures/app.py +++ b/flexmeasures/app.py @@ -89,7 +89,9 @@ def create( # Set up rate limiter app.config["RATELIMIT_STORAGE_URI"] = "redis://" - app.config["RATELIMIT_STORAGE_OPTIONS"] = {"connection_pool": redis_conn.connection_pool} + app.config["RATELIMIT_STORAGE_OPTIONS"] = { + "connection_pool": redis_conn.connection_pool + } app.limiter = Limiter( get_remote_address, app=app, From 998f19459c4e94205488fe46af2b02019f26a813 Mon Sep 17 00:00:00 2001 From: "F.N. Claessen" Date: Mon, 23 Jan 2023 10:50:07 +0100 Subject: [PATCH 3/4] pip-compile Signed-off-by: F.N. Claessen --- requirements/app.txt | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/requirements/app.txt b/requirements/app.txt index 6ecfae38f5..d4b4c4e892 100644 --- a/requirements/app.txt +++ b/requirements/app.txt @@ -38,12 +38,15 @@ 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 convertdate==2.4.0 # via workalendar cycler==0.11.0 # via matplotlib deprecated==1.2.13 # via + # limits # redis # sktime dill==0.3.5.1 @@ -64,6 +67,7 @@ flask==2.1.2 # flask-classful # flask-cors # flask-json + # flask-limiter # flask-login # flask-mail # flask-marshmallow @@ -81,6 +85,8 @@ flask-cors==3.0.10 # via -r requirements/app.in flask-json==0.3.4 # via -r requirements/app.in +flask-limiter[redis]==3.1.0 + # via -r requirements/app.in flask-login==0.5.0 # via # -r requirements/app.in @@ -119,6 +125,7 @@ idna==3.3 importlib-metadata==4.12.0 # via # -r requirements/app.in + # flask # timely-beliefs inflect==6.0.0 # via -r requirements/app.in @@ -145,6 +152,8 @@ jsonschema==4.15.0 # 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 @@ -189,8 +198,11 @@ numpy==1.22.4 # uniplot openturns==1.19.post1 # via timely-beliefs +ordered-set==4.1.0 + # via flask-limiter packaging==21.3 # via + # limits # marshmallow # marshmallow-sqlalchemy # matplotlib @@ -229,6 +241,8 @@ py-moneyed==2.0 # via -r requirements/app.in pydantic==1.10.0 # via inflect +pygments==2.14.0 + # via rich pyluach==2.0.1 # via workalendar pymeeus==0.5.11 @@ -268,6 +282,8 @@ requests==2.28.1 # tldextract requests-file==1.5.1 # via tldextract +rich==12.6.0 + # via flask-limiter rq==1.11.0 # via # -r requirements/app.in @@ -326,6 +342,8 @@ toolz==0.12.0 # via altair typing-extensions==4.3.0 # via + # flask-limiter + # limits # py-moneyed # pydantic uniplot==0.7.0 From 395c78e66db47a8363287fbfc1f58600465dbba9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicolas=20H=C3=B6ning?= Date: Thu, 2 Mar 2023 17:52:27 +0100 Subject: [PATCH 4/4] refactor: Each v3 API resource should specify their rate limits (close to the actual endpoints) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Nicolas Höning --- flexmeasures/api/v3_0/__init__.py | 27 +++++++-------------------- flexmeasures/api/v3_0/sensors.py | 13 +++++++++++-- requirements/app.txt | 1 - requirements/docs.txt | 3 ++- 4 files changed, 20 insertions(+), 24 deletions(-) diff --git a/flexmeasures/api/v3_0/__init__.py b/flexmeasures/api/v3_0/__init__.py index 2a060bbe2e..1cdb0df4a3 100644 --- a/flexmeasures/api/v3_0/__init__.py +++ b/flexmeasures/api/v3_0/__init__.py @@ -1,5 +1,4 @@ -from flask import Flask, request -from flask_security import current_user +from flask import Flask from flexmeasures.api.v3_0.sensors import SensorAPI from flexmeasures.api.v3_0.users import UserAPI @@ -12,22 +11,10 @@ def register_at(app: Flask): v3_0_api_prefix = "/api/v3_0" - def cost_function() -> int: - if request.endpoint == "SensorAPI:trigger_schedule": - return 1 - return 0 + apis = (SensorAPI, UserAPI, AssetAPI, HealthAPI) - # Apply rate limit: a schedule can be triggered once per 5 minutes per sensor per account - SensorAPI.decorators.append( - app.limiter.limit( - "1 per 5 minutes", - key_func=lambda: str(current_user.account_id) - + request.view_args.get("id", ""), - cost=cost_function, - ) - ) - - 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) + 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) diff --git a/flexmeasures/api/v3_0/sensors.py b/flexmeasures/api/v3_0/sensors.py index 6ae010b2be..704850b7a3 100644 --- a/flexmeasures/api/v3_0/sensors.py +++ b/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 @@ -54,6 +54,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( { diff --git a/requirements/app.txt b/requirements/app.txt index d4b4c4e892..38f327256a 100644 --- a/requirements/app.txt +++ b/requirements/app.txt @@ -125,7 +125,6 @@ idna==3.3 importlib-metadata==4.12.0 # via # -r requirements/app.in - # flask # timely-beliefs inflect==6.0.0 # via -r requirements/app.in diff --git a/requirements/docs.txt b/requirements/docs.txt index c2da69c97e..c720919757 100644 --- a/requirements/docs.txt +++ b/requirements/docs.txt @@ -41,8 +41,9 @@ packaging==21.3 # via # -c requirements/app.txt # sphinx -pygments==2.11.2 +pygments==2.14.0 # via + # -c requirements/app.txt # sphinx # sphinx-tabs pyparsing==3.0.9