Replace use of default dict for exemptions

Using a defaultdict results in insertions on access which can result to concorrent access errors when iteration on blueprint_exemptions. This also results in unbounded growth of both the route_exemption & blueprint_exemption dictionaries in the limit manager.
alisaifee · May 19, 2024 · 9769434 · 9769434
1 parent 7a125a6
commit 9769434
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 11 deletions.
diff --git a/flask_limiter/extension.py b/flask_limiter/extension.py
@@ -197,12 +197,8 @@ def __init__(
         self._in_memory_fallback_enabled = in_memory_fallback_enabled or (
             in_memory_fallback and len(in_memory_fallback) > 0
         )
-        self._route_exemptions: Dict[str, ExemptionScope] = defaultdict(
-            lambda: ExemptionScope.NONE
-        )
-        self._blueprint_exemptions: Dict[str, ExemptionScope] = defaultdict(
-            lambda: ExemptionScope.NONE
-        )
+        self._route_exemptions: Dict[str, ExemptionScope] = {}
+        self._blueprint_exemptions: Dict[str, ExemptionScope] = {}
         self._request_filters: List[Callable[[], bool]] = []
 
         self._headers_enabled = headers_enabled

diff --git a/flask_limiter/manager.py b/flask_limiter/manager.py
@@ -139,7 +139,7 @@ def exemption_scope(
     ) -> ExemptionScope:
         view_func = app.view_functions.get(endpoint or "", None)
         name = get_qualified_name(view_func) if view_func else ""
-        route_exemption_scope = self._route_exemptions[name]
+        route_exemption_scope = self._route_exemptions.get(name, ExemptionScope.NONE)
         blueprint_instance = app.blueprints.get(blueprint) if blueprint else None
 
         if not blueprint_instance:
@@ -161,7 +161,7 @@ def exemption_scope(
 
     def decorated_limits(self, callable_name: str) -> List[Limit]:
         limits = []
-        if not self._route_exemptions[callable_name]:
+        if not self._route_exemptions.get(callable_name, ExemptionScope.NONE):
             if callable_name in self._decorated_limits:
                 for group in self._decorated_limits[callable_name]:
                     try:
@@ -206,7 +206,9 @@ def blueprint_limits(self, app: flask.Flask, blueprint: str) -> List[Limit]:
                             limit.override_defaults for limit in blueprint_self_limits
                         )
                     )
-                    and not self._blueprint_exemptions[blueprint_name]
+                    and not self._blueprint_exemptions.get(
+                        blueprint_name, ExemptionScope.NONE
+                    )
                     & ExemptionScope.ANCESTORS
                     else blueprint_self_limits
                 )
@@ -242,7 +244,9 @@ def _blueprint_exemption_scope(
         self, app: flask.Flask, blueprint_name: str
     ) -> Tuple[ExemptionScope, Dict[str, ExemptionScope]]:
         name = app.blueprints[blueprint_name].name
-        exemption = self._blueprint_exemptions[name] & ~(ExemptionScope.ANCESTORS)
+        exemption = self._blueprint_exemptions.get(name, ExemptionScope.NONE) & ~(
+            ExemptionScope.ANCESTORS
+        )
 
         ancestory = set(blueprint_name.split("."))
         ancestor_exemption = {
@@ -251,4 +255,7 @@ def _blueprint_exemption_scope(
             if f & ExemptionScope.DESCENDENTS
         }.intersection(ancestory)
 
-        return exemption, {k: self._blueprint_exemptions[k] for k in ancestor_exemption}
+        return exemption, {
+            k: self._blueprint_exemptions.get(k, ExemptionScope.NONE)
+            for k in ancestor_exemption
+        }