diff --git a/flexmeasures/api/v3_0/__init__.py b/flexmeasures/api/v3_0/__init__.py index 56bc4d8401..1cdb0df4a3 100644 --- a/flexmeasures/api/v3_0/__init__.py +++ b/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) diff --git a/flexmeasures/api/v3_0/sensors.py b/flexmeasures/api/v3_0/sensors.py index 0b7a31b88a..7003171f65 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 @@ -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( { diff --git a/flexmeasures/api/v3_0/tests/test_sensor_schedules.py b/flexmeasures/api/v3_0/tests/test_sensor_schedules.py index fbf541419f..79b7b6ab7f 100644 --- a/flexmeasures/api/v3_0/tests/test_sensor_schedules.py +++ b/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 diff --git a/flexmeasures/app.py b/flexmeasures/app.py index 7f0a450c7f..20fb4d62a5 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,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 diff --git a/requirements/app.in b/requirements/app.in index 77005ac2ba..b325525e48 100644 --- a/requirements/app.in +++ b/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 diff --git a/requirements/app.txt b/requirements/app.txt index 9b0af92437..b00e19a198 100644 --- a/requirements/app.txt +++ b/requirements/app.txt @@ -40,6 +40,8 @@ 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 @@ -47,7 +49,10 @@ convertdate==2.4.0 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 diff --git a/requirements/docs.txt b/requirements/docs.txt index 0fb3bf1672..dfff9ce622 100644 --- a/requirements/docs.txt +++ b/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