Augment rate limiting to allow greater flexibility through the api-paste.ini configuration

1 files changed, 91 insertions, 5 deletions

diff --git a/nova/api/openstack/limits.py b/nova/api/openstack/limits.py
index d08287f6b..8642a28f9 100644
--- a/nova/api/openstack/limits.py
+++ b/nova/api/openstack/limits.py
@@ -31,8 +31,8 @@ from collections import defaultdict
 from webob.dec import wsgify
 
 from nova import quota
+from nova import utils
 from nova import wsgi as base_wsgi
-from nova import wsgi
 from nova.api.openstack import common
 from nova.api.openstack import faults
 from nova.api.openstack.views import limits as limits_views
@@ -119,6 +119,8 @@ class Limit(object):
         60 * 60 * 24: "DAY",
     }
 
+    UNIT_MAP = dict([(v, k) for k, v in UNITS.items()])
+
     def __init__(self, verb, uri, regex, value, unit):
         """
         Initialize a new `Limit`.
@@ -224,16 +226,30 @@ class RateLimitingMiddleware(base_wsgi.Middleware):
     is stored in memory for this implementation.
     """
 
-    def __init__(self, application, limits=None):
+    def __init__(self, application, limits=None, limiter=None, **kwargs):
         """
         Initialize new `RateLimitingMiddleware`, which wraps the given WSGI
         application and sets up the given limits.
 
         @param application: WSGI application to wrap
-        @param limits: List of dictionaries describing limits
+        @param limits: String describing limits
+        @param limiter: String identifying class for representing limits
+
+        Other parameters are passed to the constructor for the limiter.
         """
         base_wsgi.Middleware.__init__(self, application)
-        self._limiter = Limiter(limits or DEFAULT_LIMITS)
+
+        # Select the limiter class
+        if limiter is None:
+            limiter = Limiter
+        else:
+            limiter = utils.import_class(limiter)
+
+        # Parse the limits, if any are provided
+        if limits is not None:
+            limits = limiter.parse_limits(limits)
+
+        self._limiter = limiter(limits or DEFAULT_LIMITS, **kwargs)
 
     @wsgify(RequestClass=wsgi.Request)
     def __call__(self, req):
@@ -271,7 +287,7 @@ class Limiter(object):
     Rate-limit checking class which handles limits in memory.
     """
 
-    def __init__(self, limits):
+    def __init__(self, limits, **kwargs):
         """
         Initialize the new `Limiter`.
 
@@ -280,6 +296,12 @@ class Limiter(object):
         self.limits = copy.deepcopy(limits)
         self.levels = defaultdict(lambda: copy.deepcopy(limits))
 
+        # Pick up any per-user limit information
+        for key, value in kwargs.items():
+            if key.startswith('user:'):
+                username = key[5:]
+                self.levels[username] = self.parse_limits(value)
+
     def get_limits(self, username=None):
         """
         Return the limits for a given user.
@@ -305,6 +327,59 @@ class Limiter(object):
 
         return None, None
 
+    @staticmethod
+    def parse_limits(limits):
+        """
+        Convert a string into a list of Limit instances.  This
+        implementation expects a semicolon-separated sequence of
+        parenthesized groups, where each group contains a
+        comma-separated sequence consisting of HTTP method,
+        user-readable URI, a URI reg-exp, an integer number of
+        requests which can be made, and a unit of measure.  Valid
+        values for the latter are "SECOND", "MINUTE", "HOUR", and
+        "DAY".
+
+        @return: List of Limit instances.
+        """
+
+        # Handle empty limit strings
+        limits = limits.strip()
+        if not limits:
+            return []
+
+        # Split up the limits by semicolon
+        result = []
+        for group in limits.split(';'):
+            group = group.strip()
+            if group[0] != '(' or group[-1] != ')':
+                raise ValueError("Groups must be surrounded by parentheses")
+            group = group[1:-1]
+
+            # Extract the Limit arguments
+            args = [a.strip() for a in group.split(',')]
+            if len(args) != 5:
+                raise ValueError("Groups must contain exactly 5 elements")
+
+            # Pull out the arguments
+            verb, uri, regex, value, unit = args
+
+            # Upper-case the verb
+            verb = verb.upper()
+
+            # Convert value--raises ValueError if it's not integer
+            value = int(value)
+
+            # Convert unit
+            unit = unit.upper()
+            if unit not in Limit.UNIT_MAP:
+                raise ValueError("Invalid units specified")
+            unit = Limit.UNIT_MAP[unit]
+
+            # Build a limit
+            result.append(Limit(verb, uri, regex, value, unit))
+
+        return result
+
 
 class WsgiLimiter(object):
     """
@@ -388,3 +463,14 @@ class WsgiLimiterProxy(object):
             return None, None
 
         return resp.getheader("X-Wait-Seconds"), resp.read() or None
+
+    @staticmethod
+    def parse_limits(limits):
+        """
+        Ignore a limits string--simply doesn't apply for the limit
+        proxy.
+
+        @return: Empty list.
+        """
+
+        return []