Add Retry-After to rate-limited responses

.github/workflows/qa-tests.yml

@@ -53,4 +53,4 @@ jobs:
           dockerfile_path: ./zen-demo-python/Dockerfile
           app_port: 8080
           config_propagation_delay: 15
-          skip_tests: test-path-traversal,test-outbound-domain-blocking,test-wave-attack,test-rate-limiting-retry-after
+          skip_tests: test-path-traversal,test-outbound-domain-blocking,test-wave-attack

aikido_zen/middleware/asgi.py

@@ -19,7 +19,10 @@ async def __call__(self, scope, receive, send):
             message = "You are rate limited by Zen."
             if result["trigger"] == "ip" and result["ip"]:
                 message += " (Your IP: " + result["ip"] + ")"
-            return await send_status_code_and_text(send, (message, 429))
+            headers = {}
+            if result["retry_after_seconds"]:
+                headers["Retry-After"] = str(result["retry_after_seconds"])
+            return await send_status_code_and_text(send, (message, 429), headers)
 
         if result["type"] == "blocked":
             return await send_status_code_and_text(
@@ -30,13 +33,17 @@ async def __call__(self, scope, receive, send):
         return await self.app(scope, receive, send)
 
 
-async def send_status_code_and_text(send, pre_response):
+async def send_status_code_and_text(send, pre_response, headers=None):
     """Sends a status code and text"""
+    raw_headers = [(b"content-type", b"text/plain")]
+    for name, value in (headers or {}).items():
+        # ASGI requires lowercase header names as bytes.
+        raw_headers.append((name.lower().encode(), value.encode()))
     await send(
         {
             "type": "http.response.start",
             "status": pre_response[1],
-            "headers": [(b"content-type", b"text/plain")],
+            "headers": raw_headers,
         }
     )
     await send(

aikido_zen/middleware/asgi_test.py

@@ -0,0 +1,60 @@
+from unittest.mock import patch
+
+import pytest
+
+from .asgi import AikidoASGIMiddleware
+
+
+async def call_middleware(result):
+    sent = []
+
+    async def app(scope, receive, send):
+        raise AssertionError("app should not be called")
+
+    async def send(message):
+        sent.append(message)
+
+    with patch("aikido_zen.middleware.asgi.should_block_request", return_value=result):
+        await AikidoASGIMiddleware(app)({"type": "http"}, None, send)
+    return sent
+
+
+@pytest.mark.asyncio
+async def test_rate_limited_response_has_retry_after():
+    sent = await call_middleware(
+        {
+            "block": True,
+            "type": "ratelimited",
+            "trigger": "user",
+            "ip": "1.1.1.1",
+            "retry_after_seconds": 7,
+        }
+    )
+    assert sent[0]["status"] == 429
+    assert sent[0]["headers"] == [
+        (b"content-type", b"text/plain"),
+        (b"retry-after", b"7"),
+    ]
+    assert sent[1]["body"] == b"You are rate limited by Zen."
+
+
+@pytest.mark.asyncio
+async def test_rate_limited_response_without_retry_after_seconds():
+    sent = await call_middleware(
+        {
+            "block": True,
+            "type": "ratelimited",
+            "trigger": "user",
+            "ip": "1.1.1.1",
+            "retry_after_seconds": None,
+        }
+    )
+    assert sent[0]["status"] == 429
+    assert sent[0]["headers"] == [(b"content-type", b"text/plain")]
+
+
+@pytest.mark.asyncio
+async def test_blocked_user_response_has_no_retry_after():
+    sent = await call_middleware({"block": True, "type": "blocked", "trigger": "user"})
+    assert sent[0]["status"] == 403
+    assert sent[0]["headers"] == [(b"content-type", b"text/plain")]

aikido_zen/middleware/django.py

@@ -27,7 +27,10 @@ def __call__(self, request):
             message = "You are rate limited by Zen."
             if result["trigger"] == "ip" and result["ip"]:
                 message += " (Your IP: " + result["ip"] + ")"
-            return self.HttpResponse(message, content_type="text/plain", status=429)
+            response = self.HttpResponse(message, content_type="text/plain", status=429)
+            if result["retry_after_seconds"]:
+                response["Retry-After"] = str(result["retry_after_seconds"])
+            return response
 
         if result["type"] == "blocked":
             return self.HttpResponse(

aikido_zen/middleware/django_test.py

@@ -0,0 +1,49 @@
+from unittest.mock import patch
+
+from .django import AikidoDjangoMiddleware
+
+
+class FakeHttpResponse(dict):
+    def __init__(self, content, content_type=None, status=200):
+        super().__init__()
+        self.content = content
+        self.content_type = content_type
+        self.status_code = status
+
+
+def call_middleware(result):
+    middleware = AikidoDjangoMiddleware(get_response=None)
+    middleware.HttpResponse = FakeHttpResponse
+    with patch(
+        "aikido_zen.middleware.django.should_block_request", return_value=result
+    ):
+        return middleware(request=None)
+
+
+def test_rate_limited_response_has_retry_after():
+    response = call_middleware(
+        {
+            "block": True,
+            "type": "ratelimited",
+            "trigger": "user",
+            "ip": "1.1.1.1",
+            "retry_after_seconds": 7,
+        }
+    )
+    assert response.status_code == 429
+    assert response["Retry-After"] == "7"
+    assert response.content == "You are rate limited by Zen."
+
+
+def test_rate_limited_response_without_retry_after_seconds():
+    response = call_middleware(
+        {
+            "block": True,
+            "type": "ratelimited",
+            "trigger": "user",
+            "ip": "1.1.1.1",
+            "retry_after_seconds": None,
+        }
+    )
+    assert response.status_code == 429
+    assert "Retry-After" not in response

aikido_zen/middleware/flask.py

@@ -28,6 +28,8 @@ def __call__(self, environ, start_response):
             if result["trigger"] == "ip" and result["ip"]:
                 message += " (Your IP: " + result["ip"] + ")"
             res = self.Response(message, mimetype="text/plain", status=429)
+            if result["retry_after_seconds"]:
+                res.headers["Retry-After"] = str(result["retry_after_seconds"])
             return res(environ, start_response)
 
         if result["type"] == "blocked":

aikido_zen/middleware/flask_test.py

@@ -0,0 +1,47 @@
+from unittest.mock import patch
+
+from werkzeug.test import create_environ
+
+from .flask import AikidoFlaskMiddleware
+
+
+def call_middleware(result):
+    captured = {}
+
+    def start_response(status, headers, exc_info=None):
+        captured["status"] = status
+        captured["headers"] = dict(headers)
+
+    with patch("aikido_zen.middleware.flask.should_block_request", return_value=result):
+        middleware = AikidoFlaskMiddleware(app=None)
+        captured["body"] = b"".join(middleware(create_environ("/"), start_response))
+    return captured
+
+
+def test_rate_limited_response_has_retry_after():
+    response = call_middleware(
+        {
+            "block": True,
+            "type": "ratelimited",
+            "trigger": "ip",
+            "ip": "1.1.1.1",
+            "retry_after_seconds": 7,
+        }
+    )
+    assert response["status"] == "429 TOO MANY REQUESTS"
+    assert response["headers"]["Retry-After"] == "7"
+    assert response["body"] == b"You are rate limited by Zen. (Your IP: 1.1.1.1)"
+
+
+def test_rate_limited_response_without_retry_after_seconds():
+    response = call_middleware(
+        {
+            "block": True,
+            "type": "ratelimited",
+            "trigger": "user",
+            "ip": "1.1.1.1",
+            "retry_after_seconds": None,
+        }
+    )
+    assert response["status"] == "429 TOO MANY REQUESTS"
+    assert "Retry-After" not in response["headers"]

aikido_zen/middleware/init_test.py

@@ -161,5 +161,19 @@ def test_cache_comms_with_endpoints():
             "ip": "1.1.1.1",
             "type": "ratelimited",
             "trigger": "my_trigger",
+            "retry_after_seconds": None,
         }
         assert thread_cache.stats.rate_limited_hits == 1
+
+        mock_comms.send_data_to_bg_process.return_value = {
+            "success": True,
+            "data": {"block": True, "trigger": "my_trigger", "retry_after_seconds": 7},
+        }
+        assert should_block_request() == {
+            "block": True,
+            "ip": "1.1.1.1",
+            "type": "ratelimited",
+            "trigger": "my_trigger",
+            "retry_after_seconds": 7,
+        }
+        assert thread_cache.stats.rate_limited_hits == 2

aikido_zen/middleware/should_block_request.py

@@ -11,14 +11,13 @@
 def should_block_request():
     """
     Checks for rate-limiting and checks if the current user is blocked.
-    Returns: {
-        block: boolean,
-        data: {
-            type: "blocked" | "ratelimited",
-            trigger: "ip" | "user" | "group",
-            ip?: string
-        }
-     }
+    Returns one of:
+        {"block": False}
+        {"block": True, "type": "blocked", "trigger": "user"}
+        {"block": True, "type": "ratelimited", "trigger": "ip" | "user" | "group",
+         "ip": str | None, "retry_after_seconds": int | None}
+    retry_after_seconds is None when no retry can succeed,
+    or when the reply came from an older background process.
     """
     try:
         context = get_current_context()
@@ -69,6 +68,9 @@ def should_block_request():
                     "type": "ratelimited",
                     "trigger": ratelimit_res["data"]["trigger"],
                     "ip": context.remote_address,
+                    "retry_after_seconds": ratelimit_res["data"].get(
+                        "retry_after_seconds"
+                    ),
                 }
     except Exception as e:
         logger.debug("Exception occurred in should_block_request: %s", e)

aikido_zen/ratelimiting/__init__.py

@@ -27,37 +27,26 @@ def should_ratelimit_request(
     max_requests = int(endpoint["rateLimiting"]["maxRequests"])
     windows_size_in_ms = int(endpoint["rateLimiting"]["windowSizeInMS"])
 
+    # Only one key is checked: a group over the user, the user over the IP.
     if group:
-        allowed = connection_manager.rate_limiter.is_allowed(
-            get_key_for_group(endpoint, group),
-            windows_size_in_ms,
-            max_requests,
-        )
-        if not allowed:
-            return {"block": True, "trigger": "group"}
-
-        # Do not check IP or user rate limit if group is set
-        return {"block": False}
-    if user:
-        allowed = connection_manager.rate_limiter.is_allowed(
-            get_key_for_user(endpoint, user),
-            windows_size_in_ms,
-            max_requests,
-        )
-        if not allowed:
-            return {"block": True, "trigger": "user"}
-        # Do not check IP rate limit if user is set
+        key, trigger = get_key_for_group(endpoint, group), "group"
+    elif user:
+        key, trigger = get_key_for_user(endpoint, user), "user"
+    elif remote_address:
+        key, trigger = get_key_for_ip(endpoint, remote_address), "ip"
+    else:
         return {"block": False}
-    if remote_address:
-        allowed = connection_manager.rate_limiter.is_allowed(
-            get_key_for_ip(endpoint, remote_address),
-            windows_size_in_ms,
-            max_requests,
-        )
-        if not allowed:
-            return {"block": True, "trigger": "ip"}
 
-    return {"block": False}
+    result = connection_manager.rate_limiter.record_request(
+        key, windows_size_in_ms, max_requests
+    )
+    if result["allowed"]:
+        return {"block": False}
+    return {
+        "block": True,
+        "trigger": trigger,
+        "retry_after_seconds": result["retry_after_seconds"],
+    }
 
 
 def get_key_for_group(endpoint, group_id):