From 46dbaf95a6348482ab58d7f518fd77c1b1abb971 Mon Sep 17 00:00:00 2001
From: Dan Kruchinin <dkruchinin@acm.org>
Date: Tue, 4 Apr 2017 15:55:38 +0100
Subject: [PATCH] Response middleware should be called even if server replies
 with error

---
 sanic/app.py             | 59 +++++++++++++++++++++++-----------------
 tests/test_middleware.py | 22 +++++++++++++++
 2 files changed, 56 insertions(+), 25 deletions(-)

diff --git a/sanic/app.py b/sanic/app.py
index ef9ac57c..75e97889 100644
--- a/sanic/app.py
+++ b/sanic/app.py
@@ -444,17 +444,7 @@ class Sanic:
             # -------------------------------------------- #
 
             request.app = self
-
-            response = False
-            # The if improves speed.  I don't know why
-            if self.request_middleware:
-                for middleware in self.request_middleware:
-                    response = middleware(request)
-                    if isawaitable(response):
-                        response = await response
-                    if response:
-                        break
-
+            response = await self._run_request_middleware(request)
             # No middleware results
             if not response:
                 # -------------------------------------------- #
@@ -472,20 +462,6 @@ class Sanic:
                 response = handler(request, *args, **kwargs)
                 if isawaitable(response):
                     response = await response
-
-            # -------------------------------------------- #
-            # Response Middleware
-            # -------------------------------------------- #
-
-            if self.response_middleware:
-                for middleware in self.response_middleware:
-                    _response = middleware(request, response)
-                    if isawaitable(_response):
-                        _response = await _response
-                    if _response:
-                        response = _response
-                        break
-
         except Exception as e:
             # -------------------------------------------- #
             # Response Generation Failed
@@ -503,6 +479,17 @@ class Sanic:
                 else:
                     response = HTTPResponse(
                         "An error occurred while handling an error")
+        finally:
+            # -------------------------------------------- #
+            # Response Middleware
+            # -------------------------------------------- #
+            try:
+                response = await self._run_response_middleware(request, response)
+            except:
+                log.exception(
+                    'Exception occured in one of response middleware handlers'
+                )
+
 
         # pass the response to the correct callback
         if isinstance(response, StreamingHTTPResponse):
@@ -615,6 +602,28 @@ class Sanic:
 
         return await serve(**server_settings)
 
+    async def _run_request_middleware(self, request):
+        # The if improves speed.  I don't know why
+        if self.request_middleware:
+            for middleware in self.request_middleware:
+                response = middleware(request)
+                if isawaitable(response):
+                    response = await response
+                if response:
+                    return response
+        return None
+
+    async def _run_response_middleware(self, request, response):
+        if self.response_middleware:
+            for middleware in self.response_middleware:
+                _response = middleware(request, response)
+                if isawaitable(_response):
+                    _response = await _response
+                if _response:
+                    response = _response
+                    break
+        return response
+
     def _helper(self, host="127.0.0.1", port=8000, debug=False,
                 before_start=None, after_start=None, before_stop=None,
                 after_stop=None, ssl=None, sock=None, workers=1, loop=None,
diff --git a/tests/test_middleware.py b/tests/test_middleware.py
index f84d7151..bc1a7eb8 100644
--- a/tests/test_middleware.py
+++ b/tests/test_middleware.py
@@ -2,6 +2,7 @@ from json import loads as json_loads, dumps as json_dumps
 from sanic import Sanic
 from sanic.request import Request
 from sanic.response import json, text, HTTPResponse
+from sanic.exceptions import NotFound
 
 
 # ------------------------------------------------------------ #
@@ -53,6 +54,27 @@ def test_middleware_response():
     assert isinstance(results[2], HTTPResponse)
 
 
+def test_middleware_response_exception():
+    app = Sanic('test_middleware_response_exception')
+    result = {'status_code': None}
+
+    @app.middleware('response')
+    async def process_response(reqest, response):
+        result['status_code'] = response.status
+        return response
+
+    @app.exception(NotFound)
+    async def error_handler(request, exception):
+        return text('OK', exception.status_code)
+
+    @app.route('/')
+    async def handler(request):
+        return text('FAIL')
+
+    request, response = app.test_client.get('/page_not_found')
+    assert response.text == 'OK'
+    assert result['status_code'] == 404
+
 def test_middleware_override_request():
     app = Sanic('test_middleware_override_request')