diff --git a/samcli/local/apigw/local_apigw_service.py b/samcli/local/apigw/local_apigw_service.py index 6096085231..581e6e4507 100644 --- a/samcli/local/apigw/local_apigw_service.py +++ b/samcli/local/apigw/local_apigw_service.py @@ -152,7 +152,7 @@ def _request_handler(self, **kwargs): route.binary_types, request) except (KeyError, TypeError, ValueError): - LOG.error("Function returned an invalid response (must include one of: body, headers or " + LOG.error("Function returned an invalid response (must include one of: body, headers, multiValueHeaders or " "statusCode in the response object). Response received: %s", lambda_response) return ServiceErrorResponses.lambda_failure_response() @@ -193,8 +193,9 @@ def _parse_lambda_output(lambda_output, binary_types, flask_request): if not isinstance(json_output, dict): raise TypeError("Lambda returned %{s} instead of dict", type(json_output)) + headers = LocalApigwService._merge_response_headers(json_output.get("headers") or {}, + json_output.get("multiValueHeaders") or {}) status_code = json_output.get("statusCode") or 200 - headers = CaseInsensitiveDict(json_output.get("headers") or {}) body = json_output.get("body") or "no data" is_base_64_encoded = json_output.get("isBase64Encoded") or False @@ -234,11 +235,45 @@ def _should_base64_decode_body(binary_types, flask_request, lamba_response_heade True if the body from the request should be converted to binary, otherwise false """ - best_match_mimetype = flask_request.accept_mimetypes.best_match([lamba_response_headers["Content-Type"]]) + + # Get the first part of the content-type header, to allow for extras such as text/html; charset=utf-8 + content_type = lamba_response_headers['Content-Type'].split(";", 1)[0] + best_match_mimetype = flask_request.accept_mimetypes.best_match([content_type]) is_best_match_in_binary_types = best_match_mimetype in binary_types or '*/*' in binary_types return best_match_mimetype and is_best_match_in_binary_types and is_base_64_encoded + @staticmethod + def _merge_response_headers(headers, multi_headers): + """ + # Merge multiValueHeaders headers with headers + # Convert into CSV for Flask compatibility + # If you specify values for both headers and multiValueHeaders, API Gateway merges them into a single list. + # If the same key-value pair is specified in both, only the values from multiValueHeaders will + # appear in the merged list. + + Parameters + ---------- + headers (dict) + Headers map from the lambda_response_headers + multi_headers (dict) + multiValueHeaders map from the lambda_response_headers + + Returns + ------- + Merged list in accordance to the AWS documentation within a CaseInsensitiveDict + + """ + + processed_headers = CaseInsensitiveDict(headers) + + # Convert multi_header Lists to CSV Strings for Flask + # This gives multiValueHeaders precedence over the API Gateway headers property + for _, s in enumerate(multi_headers): + processed_headers[s] = ", ".join(multi_headers[s]) + + return processed_headers + @staticmethod def _construct_event(flask_request, port, binary_types): """ diff --git a/tests/unit/local/apigw/test_local_apigw_service.py b/tests/unit/local/apigw/test_local_apigw_service.py index 234554deca..54771be238 100644 --- a/tests/unit/local/apigw/test_local_apigw_service.py +++ b/tests/unit/local/apigw/test_local_apigw_service.py @@ -260,6 +260,47 @@ def test_class_initialization(self): self.assertEquals(self.api_gateway.path, '/') +class TestLambdaHeaderDictionaryMerge(TestCase): + def test_empty_dictionaries_produce_empty_result(self): + headers = {} + multi_value_headers = {} + + result = LocalApigwService._merge_response_headers(headers, multi_value_headers) + + self.assertEquals(result, {}) + + def test_headers_are_merged(self): + headers = {"h1": "value1", "h2": "value2"} + multi_value_headers = {"h3": ["value3"]} + + result = LocalApigwService._merge_response_headers(headers, multi_value_headers) + + self.assertIn("h1", result) + self.assertIn("h2", result) + self.assertIn("h3", result) + self.assertEquals(result["h1"], "value1") + self.assertEquals(result["h2"], "value2") + self.assertEquals(result["h3"], "value3") + + def test_multivalue_headers_are_turned_into_csv(self): + headers = {} + multi_value_headers = {"h1": ["a", "b", "c"]} + + result = LocalApigwService._merge_response_headers(headers, multi_value_headers) + + self.assertIn("h1", result) + self.assertEquals(result["h1"], "a, b, c") + + def test_multivalue_headers_override_headers_dict(self): + headers = {"h1": "ValueA"} + multi_value_headers = {"h1": ["ValueB"]} + + result = LocalApigwService._merge_response_headers(headers, multi_value_headers) + + self.assertIn("h1", result) + self.assertEquals(result["h1"], "ValueB") + + class TestServiceParsingLambdaOutput(TestCase): def test_default_content_type_header_added_with_no_headers(self): @@ -289,6 +330,15 @@ def test_custom_content_type_header_is_not_modified(self): self.assertIn("Content-Type", headers) self.assertEquals(headers["Content-Type"], "text/xml") + def test_custom_content_type_multivalue_header_is_not_modified(self): + lambda_output = '{"statusCode": 200, "multiValueHeaders":{"Content-Type": ["text/xml"]}, "body": "{}", ' \ + '"isBase64Encoded": false}' + + (_, headers, _) = LocalApigwService._parse_lambda_output(lambda_output, binary_types=[], flask_request=Mock()) + + self.assertIn("Content-Type", headers) + self.assertEquals(headers["Content-Type"], "text/xml") + def test_extra_values_ignored(self): lambda_output = '{"statusCode": 200, "headers": {}, "body": "{\\"message\\":\\"Hello from Lambda\\"}", ' \ '"isBase64Encoded": false, "another_key": "some value"}'