Fix for GH#107: Enable support for multiple Set-Cookie response headers (native & wsgi/asgi) (#115)

* initial multi cookie attempt

* address pr comments

Co-authored-by: peterstone2017 <[email protected]>

Author: YunchuWang

azure/functions/_abc.py

@@ -6,6 +6,7 @@
 import io
 import typing
 
+from azure.functions._thirdparty.werkzeug.datastructures import Headers
 
 T = typing.TypeVar('T')
 
@@ -193,7 +194,7 @@ def charset(self):
 
     @property
     @abc.abstractmethod
-    def headers(self) -> typing.MutableMapping[str, str]:
+    def headers(self) -> Headers:
         pass
 
     @abc.abstractmethod

azure/functions/_http.py

@@ -4,14 +4,14 @@
 import collections.abc
 import io
 import json
-import typing
 import types
+import typing
 
 from . import _abc
-
 from ._thirdparty.werkzeug import datastructures as _wk_datastructures
 from ._thirdparty.werkzeug import formparser as _wk_parser
 from ._thirdparty.werkzeug import http as _wk_http
+from ._thirdparty.werkzeug.datastructures import Headers
 
 
 class BaseHeaders(collections.abc.Mapping):
@@ -40,13 +40,8 @@ class HttpRequestHeaders(BaseHeaders):
     pass
 
 
-class HttpResponseHeaders(BaseHeaders, collections.abc.MutableMapping):
-
-    def __setitem__(self, key: str, value: str):
-        self.__http_headers__[key.lower()] = value
-
-    def __delitem__(self, key: str):
-        del self.__http_headers__[key.lower()]
+class HttpResponseHeaders(Headers):
+    pass
 
 
 class HttpResponse(_abc.HttpResponse):
@@ -90,7 +85,10 @@ def __init__(self,
 
         if headers is None:
             headers = {}
-        self.__headers = HttpResponseHeaders(headers)
+
+        self.__headers = HttpResponseHeaders([])
+        for k, v in headers.items():
+            self.__headers.add_header(k, v)
 
         if body is not None:
             self.__set_body(body)

azure/functions/http.py

@@ -3,11 +3,12 @@
 
 import json
 import typing
+from http.cookies import SimpleCookie
 
 from azure.functions import _abc as azf_abc
 from azure.functions import _http as azf_http
-
 from . import meta
+from ._thirdparty.werkzeug.datastructures import Headers
 
 
 class HttpRequest(azf_http.HttpRequest):
@@ -79,7 +80,8 @@ def encode(cls, obj: typing.Any, *,
 
         if isinstance(obj, azf_abc.HttpResponse):
             status = obj.status_code
-            headers = dict(obj.headers)
+            headers: Headers = obj.headers
+
             if 'content-type' not in headers:
                 if obj.mimetype.startswith('text/'):
                     ct = f'{obj.mimetype}; charset={obj.charset}'
@@ -93,6 +95,12 @@ def encode(cls, obj: typing.Any, *,
             else:
                 datum_body = meta.Datum(type='bytes', value=b'')
 
+            cookies = None
+            if "Set-Cookie" in headers:
+                cookies = [SimpleCookie(cookie) for cookie in
+                           headers.get_all('Set-Cookie')]
+                headers.pop("Set-Cookie")
+
             return meta.Datum(
                 type='http',
                 value=dict(
@@ -101,6 +109,7 @@ def encode(cls, obj: typing.Any, *,
                         n: meta.Datum(type='string', value=h)
                         for n, h in headers.items()
                     },
+                    cookies=cookies,
                     body=datum_body,
                 )
             )

tests/test_http.py

@@ -5,6 +5,7 @@
 
 import azure.functions as func
 import azure.functions.http as http
+from azure.functions._http import HttpResponseHeaders
 
 
 class TestHTTP(unittest.TestCase):
@@ -88,6 +89,53 @@ def test_http_output_type(self):
         self.assertTrue(check_output_type(func.HttpResponse))
         self.assertTrue(check_output_type(str))
 
+    def test_http_response_encode_to_datum_no_cookie(self):
+        resp = func.HttpResponse()
+        datum = http.HttpResponseConverter.encode(resp, expected_type=None)
+
+        self.assertEqual(datum.value["cookies"], None)
+
+    def test_http_response_encode_to_datum_with_cookies(self):
+        headers = HttpResponseHeaders()
+        headers.add("Set-Cookie",
+                    'foo3=42; Domain=example.com; Expires=Thu, '
+                    '12-Jan-2017 13:55:08 GMT; Path=/; Max-Age=10000000')
+        headers.add("Set-Cookie",
+                    'foo3=43; Domain=example.com; Expires=Thu, 12-Jan-2018 '
+                    '13:55:09 GMT; Path=/; Max-Age=10000000')
+        resp = func.HttpResponse(headers=headers)
+        datum = http.HttpResponseConverter.encode(resp, expected_type=None)
+
+        actual_cookies = datum.value['cookies']
+        self.assertIsNotNone(actual_cookies)
+        self.assertTrue(isinstance(actual_cookies, list))
+        self.assertTrue(len(actual_cookies), 2)
+        self.assertEqual(str(actual_cookies[0]),
+                         "Set-Cookie: foo3=42; Domain=example.com; "
+                         "expires=Thu, 12-Jan-2017 13:55:08 GMT; "
+                         "Max-Age=10000000; Path=/")
+        self.assertEqual(str(actual_cookies[1]),
+                         "Set-Cookie: foo3=43; Domain=example.com; "
+                         "expires=Thu, 12-Jan-2018 13:55:09 GMT; "
+                         "Max-Age=10000000; Path=/")
+
+        self.assertTrue("Set-Cookie" not in resp.headers)
+
+    def test_http_response_encode_to_datum_with_cookies_lower_case(self):
+        headers = HttpResponseHeaders()
+        headers.add("set-cookie",
+                    'foo3=42; Domain=example.com; Path=/; Max-Age=10000.0')
+        resp = func.HttpResponse(headers=headers)
+        datum = http.HttpResponseConverter.encode(resp, expected_type=None)
+
+        actual_cookies = datum.value['cookies']
+        self.assertIsNotNone(actual_cookies)
+        self.assertTrue(isinstance(actual_cookies, list))
+        self.assertTrue(len(actual_cookies), 1)
+        self.assertEqual(str(actual_cookies[0]),
+                         "Set-Cookie: foo3=42; Domain=example.com; "
+                         "Max-Age=10000.0; Path=/")
+
     def test_http_request_should_not_have_implicit_output(self):
         self.assertFalse(http.HttpRequestConverter.has_implicit_output())
 

tests/test_http_wsgi.py

@@ -6,6 +6,7 @@
 
 import azure.functions as func
 from azure.functions._abc import TraceContext, RetryContext
+from azure.functions._http import HttpResponseHeaders
 from azure.functions._http_wsgi import (
     WsgiRequest,
     WsgiResponse,
@@ -153,7 +154,7 @@ def test_response_no_headers(self):
 
         wsgi_response: WsgiResponse = WsgiResponse.from_app(app, environ)
         func_response: func.HttpResponse = wsgi_response.to_func_response()
-        self.assertEqual(func_response.headers, {})
+        self.assertEqual(func_response.headers, HttpResponseHeaders([]))
 
     def test_response_with_exception(self):
         app = self._generate_wsgi_app(