Added a pre-processing decorator at api.response.message_extractor

This is to aid custom ``response_callback`` implementation for use in
application bindings like ``wsgi.InterceptorMiddleware`` and
``cgi.InterceptorController``. It simplifies the need to pre-process
response body data and ensures consistency through hiding details.

Author: liuyu81

src/SnapSearch/api/response.py

@@ -12,9 +12,11 @@
     :date: 2014/03/08
 """
 
-__all__ = ['Response', ]
+__all__ = ['Response', 'message_extractor', ]
 
 
+import functools
+
 from .._compat import u, HTTP_STATUS_CODES
 
 
@@ -28,6 +30,8 @@ def _extract_message(response_body):
     payload = response_body.get('html', u("")).encode("utf-8")
 
     # response headers
+    found_content_length = False
+
     headers = []
     for item in response_body.get('headers', []):
 
@@ -44,19 +48,32 @@ def _extract_message(response_body):
                str(item['value']).encode("utf-8"))
 
         # dirty entry (needs to be fixed later)
-        if tup[0] in (b"content-type", b"content-length"):
+        if tup[0] == b"content-length":
+            found_content_length = True
             continue
 
         # selected entry
         headers.append(tup)
 
-    # fixed header entries
-    headers.append((b"content-type", b"text/html"))
-    headers.append((b"content-length", bytes(len(payload))))
+    # fixed content-length to avoid truncation of content by the web server in
+    # case the headers are to be replicated in an HTTP response message.
+    if found_content_length:
+        headers.append((b"content-length", bytes(len(payload))))
 
     return {'status': status, 'headers': headers, 'html': payload}
 
 
+def message_extractor(func):
+    """
+    Decorator for data extracting functions requiring the HTTP message fields
+    from a raw response body. This is an input preprocessing decorator for the
+    implementation of ``response_callback`` used in application bindings.
+    """
+    def wrapper(response_body):
+        return func(_extract_message(response_body))
+    return functools.update_wrapper(wrapper, func)
+
+
 class Response(object):
     """
     Wraps an HTTP response from SnapSearch's backend service.

src/SnapSearch/cgi.py

@@ -22,6 +22,17 @@
 from .interceptor import Interceptor
 
 
+@api.response.message_extractor
+def default_response_callback(response_body):
+    """
+    Removes staled HTTP headers from response body
+    """
+    response_body['headers'] = [
+        (key, val) for key, val in response_body['headers']
+        if key.lower() in (b"location", b"server", b"status")]
+    return response_body
+
+
 class InterceptorController(object):
     """
     Wraps a CGI script (see :RFC:`3875`) by temporarily buffering its standard
@@ -71,7 +82,7 @@ def __init__(self, interceptor, response_callback=None):
         #
         self.__interceptor = interceptor
         self.__response_callback = response_callback \
-            if callable(response_callback) else api.response._extract_message
+            if callable(response_callback) else default_response_callback
         self.__real_stdout = None
         self.__stdout_buffer = None
         pass

src/SnapSearch/wsgi.py

@@ -21,6 +21,17 @@
 from .interceptor import Interceptor
 
 
+@api.response.message_extractor
+def default_response_callback(response_body):
+    """
+    Removes staled HTTP headers from response body
+    """
+    response_body['headers'] = [
+        (key, val) for key, val in response_body['headers']
+        if key.lower() in (b"location", b"server", b"status")]
+    return response_body
+
+
 class InterceptorMiddleware(object):
     """
     Wraps a WSGI-defined web application (see :PEP:`3333`) and intercepts
@@ -77,7 +88,7 @@ def __init__(self, application, interceptor, response_callback=None):
         self.__application = application
         self.__interceptor = interceptor
         self.__response_callback = response_callback \
-            if callable(response_callback) else api.response._extract_message
+            if callable(response_callback) else default_response_callback
         pass
 
     def __call__(self, environ, start_response):

tests/test_cgi.py

@@ -62,7 +62,8 @@ def test_controller_init(self):
         from SnapSearch.cgi import InterceptorController
         i = Interceptor(self.client, self.detector)
         cb = lambda r: {b"status": 200,
-                        b"headers": [(b"Content-Type", b"text/html"), ],
+                        b"headers": [(b"Content-Type", b"text/html"),
+                                     (b"Server", b"apache (CentOS)"), ],
                         b"html": b"Hello World!\r\n", }
         ic = InterceptorController(i, cb)
         self.assertEqual(i, ic.interceptor)

tests/test_wsgi.py

@@ -74,7 +74,8 @@ def test_middleware_init(self):
         a = DummyApp()
         i = Interceptor(self.client, self.detector)
         cb = lambda r: {b"status": 200,
-                        b"headers": [(b"Content-Type", b"text/html"), ],
+                        b"headers": [(b"Content-Type", b"text/html"),
+                                     (b"Server", b"apache (CentOS)"), ],
                         b"html": b"Hello World!\r\n", }
         im = InterceptorMiddleware(a, i, cb)
         self.assertEqual(a, im.application)