Rate limit marketo endpoint

tests/test_decorators.py

@@ -0,0 +1,48 @@
+# Packages
+from time import sleep
+from unittest import TestCase
+
+import werkzeug
+
+from webapp.app import app
+from webapp.decorators import rate_limit_with_backoff
+
+
+class TestDecorators(TestCase):
+    def test_rate_limit_with_backoff_blocks_requests(self):
+        """
+        Test that functions generated by rate_limit_with_backoff are rate
+        limited and that they backoff when the rate limit is exceeded.
+        """
+
+        def fn():
+            sleep(0.1)
+            return True
+
+        with app.test_request_context():
+            # Limit to calls once every second
+            rate_limited_fn = rate_limit_with_backoff(fn, (1, 1))
+
+            # Should raise an exception
+            with self.assertRaises(werkzeug.exceptions.TooManyRequests):
+                while True:
+                    rate_limited_fn()
+
+    def test_rate_limit_with_backoff_allows_requests(self):
+        """
+        Test that functions generated by rate_limit_with_backoff are rate
+        limited and that they backoff when the rate limit is exceeded.
+        """
+
+        def fn():
+            sleep(0.1)
+            return True
+
+        with app.test_request_context():
+            # Limit to calls once every second
+            rate_limited_fn = rate_limit_with_backoff(fn, (1, 1))
+
+            # Should not raise an exception
+            for _ in range(3):
+                sleep(1)
+                rate_limited_fn()

webapp/decorators.py

@@ -1,8 +1,12 @@
 # Core packages
 import functools
+import json
+from datetime import datetime, timedelta
+from typing import Callable, Optional
 
 # Third party packages
 import flask
+
 from webapp.login import user_info
 
 
@@ -20,3 +24,67 @@ def is_user_logged_in(*args, **kwargs):
         return func(*args, **kwargs)
 
     return is_user_logged_in
+
+
+def rate_limit_with_backoff(
+    func: Callable, limits: Optional[tuple[int, int]] = None
+) -> Callable:
+    """
+    Decorator to rate limit function calls based on the users'
+    session. The default rate limit restricts users to:
+    - 1 request every 4 seconds
+    - 4 requests every 60 seconds
+
+    This can be overwritten with the limits argument e.g.
+        @rate_limit_with_backoff(limits=(1, 10))
+    for 1 request every 10 seconds.
+
+    @param func: Function to decorate
+    @param limits: Tuple of (requests, seconds) request limit mappings
+    """
+
+    rate_limit_attempt_map = {
+        1: timedelta(seconds=4),
+        4: timedelta(seconds=16),
+        16: timedelta(seconds=64),
+    }
+
+    if limits:
+        additional_limits = {limits[0]: timedelta(seconds=limits[1])}
+        rate_limit_attempt_map = additional_limits
+
+    @functools.wraps(func)
+    def rate_limited(*args, **kwargs):
+        try:
+            # Get the initial request
+            initial_request = json.loads(flask.session[func.__name__])
+            for limit in sorted(rate_limit_attempt_map.keys()):
+                # Get the seconds limit for these attempts
+                if limit >= initial_request["attempts"]:
+                    seconds_limit = rate_limit_attempt_map.get(limit)
+                    time_since_last_request = (
+                        datetime.now()
+                        - datetime.fromtimestamp(initial_request["timestamp"])
+                    )
+                    # Abort if the request is too soon
+                    if (
+                        time_since_last_request.total_seconds()
+                        < seconds_limit.total_seconds()
+                    ):
+                        return flask.abort(429)
+                    break
+
+            # Reset the timestamp if the request succeeds
+            initial_request["timestamp"] = datetime.now()
+
+            # Otherwise update the session
+            initial_request["attempts"] += 1
+            flask.session[func.__name__] = json.dumps(initial_request)
+        except (KeyError, TypeError):
+            # Set values for initial request
+            flask.session[func.__name__] = json.dumps(
+                {"timestamp": datetime.now().timestamp(), "attempts": 1}
+            )
+        return func(*args, **kwargs)
+
+    return rate_limited

webapp/views.py

@@ -1,35 +1,35 @@
 # Standard library
 import datetime
 import html
+import json
 import math
 import os
 import re
-import json
+from urllib.parse import quote, unquote
 
 # Packages
 import dateutil
 import feedparser
 import flask
+import jinja2
 import talisker.requests
 import yaml
-import jinja2
-from ubuntu_release_info.data import Data
-from geolite2 import geolite2
-from requests import Session
-from requests.exceptions import HTTPError
-from urllib.parse import quote, unquote
-
-from canonicalwebteam.search.models import get_search_results
-from canonicalwebteam.search.views import NoAPIKeyError
 from bs4 import BeautifulSoup
-from werkzeug.exceptions import BadRequest
 from canonicalwebteam.discourse import (
     DiscourseAPI,
-    Docs,
     DocParser,
+    Docs,
 )
+from canonicalwebteam.search.models import get_search_results
+from canonicalwebteam.search.views import NoAPIKeyError
+from geolite2 import geolite2
+from requests import Session
+from requests.exceptions import HTTPError
+from ubuntu_release_info.data import Data
+from werkzeug.exceptions import BadRequest
 
 # Local
+from webapp.decorators import rate_limit_with_backoff
 from webapp.login import user_info
 from webapp.marketo import MarketoAPI
 
@@ -899,6 +899,7 @@ def shorten_acquisition_url(acquisition_url):
     return acquisition_url
 
 
+@rate_limit_with_backoff
 def marketo_submit():
     form_fields = {}
     for key in flask.request.form: