middlewares: support multiple start_response calls

Previously middlewares mostly gave up and acted as pass through when a
second start_response call occurred. Now they try to handle this
situation.

1 files changed, 34 insertions, 28 deletions

diff --git a/wsgitools/middlewares.py b/wsgitools/middlewares.py
index fa2de2e..9792089 100644
--- a/wsgitools/middlewares.py
+++ b/wsgitools/middlewares.py
@@ -143,45 +143,49 @@ class ContentLengthMiddleware:
         """wsgi interface"""
         assert isinstance(environ, dict)
         todo = []
+        gotdata = False
         def modified_start_response(status, headers, exc_info=None):
+            try:
+                if gotdata:
+                    assert exc_info is not None
+                    raise exc_info[0], exc_info[1], exc_info[2]
+            finally:
+                exc_info = None
             assert isinstance(status, str)
             assert isinstance(headers, list)
-            if (exc_info is not None or
-                 [v for h, v in headers if h.lower() == "content-length"]):
-                todo[:] = (None,)
-                return start_response(status, headers, exc_info)
-            else:
-                todo[:] = ((status, headers),)
-                def raise_not_imp(*args):
-                    raise NotImplementedError
-                return raise_not_imp
+            todo[:] = ((status, headers),)
+            def raise_not_imp(*args):
+                raise NotImplementedError
+            return raise_not_imp
 
         ret = self.app(environ, modified_start_response)
         assert hasattr(ret, "__iter__")
 
-        if todo and todo[0] is None: # nothing to do
-            #print "content-length: nothing"
-            return ret
-
         if isinstance(ret, list):
-            #print "content-length: simple"
+            gotdata = True
+            assert bool(todo)
             status, headers = todo[0]
-            length = sum(map(len, ret))
-            headers.append(("Content-length", str(length)))
+            if not [k for k, _ in headers if k.lower() == "content-length"]:
+                length = sum(map(len, ret))
+                headers.append(("Content-Length", str(length)))
             start_response(status, headers)
             return ret
 
         ret = iter(ret)
+        first = ""
         stopped = False
-        data = CloseableList(getattr(ret, "close", None))
-        length = 0
-        try:
-            data.append(next(ret)) # fills todo
-            length += len(data[-1])
-        except StopIteration:
-            stopped = True
-
+        while not (first or stopped):
+            try:
+                first = next(ret)
+            except StopIteration:
+                stopped = True
+        gotdata = True
+        assert bool(todo)
         status, headers = todo[0]
+        data = CloseableList(getattr(ret, "close", None))
+        if first:
+            data.append(first)
+        length = len(first)
 
         while (not stopped) and length < self.maxstore:
             try:
@@ -191,12 +195,10 @@ class ContentLengthMiddleware:
                 stopped = True
 
         if stopped:
-            #print "content-length: gathered"
             headers.append(("Content-length", str(length)))
             start_response(status, headers)
             return data
 
-        #print "content-length: passthrough"
         start_response(status, headers)
 
         return CloseableIterator(getattr(ret, "close", None), data, ret)
@@ -276,10 +278,14 @@ class CachingMiddleware:
                 del self.cache[path]
         cache_object = [now, "", [], []]
         def modified_start_respesponse(status, headers, exc_info=None):
+            try:
+                if cache_object[3]:
+                    assert exc_info is not None
+                    raise exc_info[0], exc_info[1], exc_info[2]
+            finally:
+                exc_info = None
             assert isinstance(status, str)
             assert isinstance(headers, list)
-            if exc_info is not None:
-                return self.app(status, headers, exc_info)
             cache_object[1] = status
             cache_object[2] = headers
             write = start_response(status, list(headers))