Merge "Updates to OSAPI sizelimit middleware."

1 files changed, 36 insertions, 5 deletions

diff --git a/nova/api/sizelimit.py b/nova/api/sizelimit.py
index 67d459583..1e88f183e 100644
--- a/nova/api/sizelimit.py
+++ b/nova/api/sizelimit.py
@@ -37,6 +37,35 @@ CONF.register_opt(max_request_body_size_opt)
 LOG = logging.getLogger(__name__)
 
 
+class LimitingReader(object):
+    """Reader to limit the size of an incoming request."""
+    def __init__(self, data, limit):
+        """
+        :param data: Underlying data object
+        :param limit: maximum number of bytes the reader should allow
+        """
+        self.data = data
+        self.limit = limit
+        self.bytes_read = 0
+
+    def __iter__(self):
+        for chunk in self.data:
+            self.bytes_read += len(chunk)
+            if self.bytes_read > self.limit:
+                msg = _("Request is too large.")
+                raise webob.exc.HTTPRequestEntityTooLarge(explanation=msg)
+            else:
+                yield chunk
+
+    def read(self, i=None):
+        result = self.data.read(i)
+        self.bytes_read += len(result)
+        if self.bytes_read > self.limit:
+            msg = _("Request is too large.")
+            raise webob.exc.HTTPRequestEntityTooLarge(explanation=msg)
+        return result
+
+
 class RequestBodySizeLimiter(wsgi.Middleware):
     """Limit the size of incoming requests."""
 
@@ -45,9 +74,11 @@ class RequestBodySizeLimiter(wsgi.Middleware):
 
     @webob.dec.wsgify(RequestClass=wsgi.Request)
     def __call__(self, req):
-        if (req.content_length > CONF.osapi_max_request_body_size
-            or len(req.body) > CONF.osapi_max_request_body_size):
+        if req.content_length > CONF.osapi_max_request_body_size:
             msg = _("Request is too large.")
-            raise webob.exc.HTTPBadRequest(explanation=msg)
-        else:
-            return self.application
+            raise webob.exc.HTTPRequestEntityTooLarge(explanation=msg)
+        if req.content_length is None and req.is_body_readable:
+            limiter = LimitingReader(req.body_file,
+                                     CONF.osapi_max_request_body_size)
+            req.body_file = limiter
+        return self.application