allow for existing retry after header

flask_limiter/extension.py

@@ -1,6 +1,7 @@
 """
 the flask extension
 """
+import datetime
 import itertools
 import logging
 import sys
@@ -13,7 +14,7 @@
 from limits.errors import ConfigurationError
 from limits.storage import storage_from_string, MemoryStorage
 from limits.strategies import STRATEGIES
-from werkzeug.http import http_date
+from werkzeug.http import http_date, parse_date
 
 from flask_limiter.wrappers import Limit, LimitGroup
 from .errors import RateLimitExceeded
@@ -310,7 +311,25 @@ def __inject_headers(self, response):
                 self._header_mapping[HEADERS.REMAINING], window_stats[1]
             )
             response.headers.add(self._header_mapping[HEADERS.RESET], reset_in)
-            response.headers.add(
+
+            # response may have an existing retry after
+            existing_retry_after_header = response.headers.get('Retry-After')
+
+            if existing_retry_after_header is not None:
+                # might be in http-date format
+                retry_after = parse_date(existing_retry_after_header)
+
+                # parse_date failure returns None
+                if retry_after is None:
+                    retry_after = time.time() + int(existing_retry_after_header)
+
+                if isinstance(retry_after, datetime.datetime):
+                    retry_after = time.mktime(retry_after.timetuple())
+
+                reset_in = max(retry_after, reset_in)
+
+            # set the header instead of using add
+            response.headers.set(
                 self._header_mapping[HEADERS.RETRY_AFTER],
                 self._retry_after == 'http-date' and http_date(reset_in)
                 or int(reset_in - time.time())

tests/test_flask_ext.py

@@ -1287,6 +1287,44 @@ def t():
                 resp = cli.get("/t1")
                 self.assertEqual(resp.status_code, 200)
 
+    def test_retry_after_exists_seconds(self):
+        app = Flask(__name__)
+        _ = Limiter(
+            app,
+            default_limits=["1/minute"],
+            headers_enabled=True,
+            key_func=get_remote_address
+        )
+
+        @app.route("/t1")
+        def t():
+            return "", 200, {'Retry-After': '1000000'}
+
+        with app.test_client() as cli:
+            resp = cli.get("/t1")
+
+            retry_after = int(resp.headers.get('Retry-After'))
+            self.assertTrue(retry_after > 1000)
+
+    def test_retry_after_exists_rfc1123(self):
+        app = Flask(__name__)
+        _ = Limiter(
+            app,
+            default_limits=["1/minute"],
+            headers_enabled=True,
+            key_func=get_remote_address
+        )
+
+        @app.route("/t1")
+        def t():
+            return "", 200, {'Retry-After': 'Sun, 06 Nov 2032 01:01:01 GMT'}
+
+        with app.test_client() as cli:
+            resp = cli.get("/t1")
+
+            retry_after = int(resp.headers.get('Retry-After'))
+            self.assertTrue(retry_after > 1000)
+
     def test_custom_headers_from_setter(self):
         app = Flask(__name__)
         limiter = Limiter(