Fixed #16004 - csrf_protect does not send cookie if view returns TemplateResponse

The root bug was in decorator_from_middleware, and the fix also corrects
bugs with gzip_page and other decorators.

git-svn-id: http://code.djangoproject.com/svn/django/trunk@16276 bcc190cf-cafb-0310-a4f2-bffc1f526a37

5 files changed, 88 insertions, 7 deletions

diff --git a/django/core/handlers/base.py b/django/core/handlers/base.py
index f216886d56..d653860547 100644
--- a/django/core/handlers/base.py
+++ b/django/core/handlers/base.py
@@ -133,7 +133,7 @@ class BaseHandler(object):
                 if hasattr(response, 'render') and callable(response.render):
                     for middleware_method in self._template_response_middleware:
                         response = middleware_method(request, response)
-                    response.render()
+                    response = response.render()
 
             except http.Http404, e:
                 logger.warning('Not Found: %s' % request.path,
diff --git a/django/template/response.py b/django/template/response.py
index a6ef893520..73645a7d72 100644
--- a/django/template/response.py
+++ b/django/template/response.py
@@ -92,11 +92,14 @@ class SimpleTemplateResponse(HttpResponse):
 
         Returns the baked response instance.
         """
+        retval = self
         if not self._is_rendered:
             self._set_content(self.rendered_content)
             for post_callback in self._post_render_callbacks:
-                post_callback(self)
-        return self
+                newretval = post_callback(retval)
+                if newretval is not None:
+                    retval = newretval
+        return retval
 
     is_rendered = property(lambda self: self._is_rendered)
 
diff --git a/django/utils/decorators.py b/django/utils/decorators.py
index ee169ec7cc..22f33a76a4 100644
--- a/django/utils/decorators.py
+++ b/django/utils/decorators.py
@@ -95,10 +95,17 @@ def make_middleware_decorator(middleware_class):
                         if result is not None:
                             return result
                     raise
-                if hasattr(middleware, 'process_response'):
-                    result = middleware.process_response(request, response)
-                    if result is not None:
-                        return result
+                if hasattr(response, 'render') and callable(response.render):
+                    if hasattr(middleware, 'process_template_response'):
+                        response = middleware.process_template_response(request, response)
+                    # Defer running of process_response until after the template
+                    # has been rendered:
+                    if hasattr(middleware, 'process_response'):
+                        callback = lambda response: middleware.process_response(request, response)
+                        response.add_post_render_callback(callback)
+                else:
+                    if hasattr(middleware, 'process_response'):
+                        return middleware.process_response(request, response)
                 return response
             return _wrapped_view
         return _decorator
diff --git a/docs/ref/template-response.txt b/docs/ref/template-response.txt
index 1b92d73d48..0064290d4e 100644
--- a/docs/ref/template-response.txt
+++ b/docs/ref/template-response.txt
@@ -119,6 +119,10 @@ Methods
     rendered :class:`~django.template.response.SimpleTemplateResponse`
     instance.
 
+    If the callback returns a value that is not `None`, this will be
+    used as the response instead of the original response object (and
+    will be passed to the next post rendering callback etc.)
+
 .. method:: SimpleTemplateResponse.render():
 
     Sets :attr:`response.content` to the result obtained by
diff --git a/tests/regressiontests/utils/decorators.py b/tests/regressiontests/utils/decorators.py
index db32418de7..837cc7d39e 100644
--- a/tests/regressiontests/utils/decorators.py
+++ b/tests/regressiontests/utils/decorators.py
@@ -1,5 +1,7 @@
 from django.http import HttpResponse
 from django.middleware.doc import XViewMiddleware
+from django.template import Template, Context
+from django.template.response import TemplateResponse
 from django.test import TestCase, RequestFactory
 from django.utils.decorators import decorator_from_middleware
 
@@ -19,6 +21,26 @@ class ClassXView(object):
 class_xview = xview_dec(ClassXView())
 
 
+class FullMiddleware(object):
+    def process_request(self, request):
+        request.process_request_reached = True
+
+    def process_view(sef, request, view_func, view_args, view_kwargs):
+        request.process_view_reached = True
+
+    def process_template_response(self, request, response):
+        request.process_template_response_reached = True
+        return response
+
+    def process_response(self, request, response):
+        # This should never receive unrendered content.
+        request.process_response_content = response.content
+        request.process_response_reached = True
+        return response
+
+full_dec = decorator_from_middleware(FullMiddleware)
+
+
 class DecoratorFromMiddlewareTests(TestCase):
     """
     Tests for view decorators created using
@@ -37,3 +59,48 @@ class DecoratorFromMiddlewareTests(TestCase):
         Test a middleware that implements process_view, operating on a callable class.
         """
         class_xview(self.rf.get('/'))
+
+    def test_full_dec_normal(self):
+        """
+        Test that all methods of middleware are called for normal HttpResponses
+        """
+
+        @full_dec
+        def normal_view(request):
+            t = Template("Hello world")
+            return HttpResponse(t.render(Context({})))
+
+        request = self.rf.get('/')
+        response = normal_view(request)
+        self.assertTrue(getattr(request, 'process_request_reached', False))
+        self.assertTrue(getattr(request, 'process_view_reached', False))
+        # process_template_response must not be called for HttpResponse
+        self.assertFalse(getattr(request, 'process_template_response_reached', False))
+        self.assertTrue(getattr(request, 'process_response_reached', False))
+
+    def test_full_dec_templateresponse(self):
+        """
+        Test that all methods of middleware are called for TemplateResponses in
+        the right sequence.
+        """
+
+        @full_dec
+        def template_response_view(request):
+            t = Template("Hello world")
+            return TemplateResponse(request, t, {})
+
+        request = self.rf.get('/')
+        response = template_response_view(request)
+        self.assertTrue(getattr(request, 'process_request_reached', False))
+        self.assertTrue(getattr(request, 'process_view_reached', False))
+        self.assertTrue(getattr(request, 'process_template_response_reached', False))
+        # response must not be rendered yet.
+        self.assertFalse(response._is_rendered)
+        # process_response must not be called until after response is rendered,
+        # otherwise some decorators like csrf_protect and gzip_page will not
+        # work correctly. See #16004
+        self.assertFalse(getattr(request, 'process_response_reached', False))
+        response.render()
+        self.assertTrue(getattr(request, 'process_response_reached', False))
+        # Check that process_response saw the rendered content
+        self.assertEqual(request.process_response_content, "Hello world")