Add exposed headers to gateway responses

smithy-lang · Jul 6, 2021 · b21eb7b · b21eb7b
1 parent 2b90a8e
commit b21eb7b
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 9 deletions.
diff --git a/...rc/main/java/software/amazon/smithy/aws/apigateway/openapi/AddCorsToGatewayResponses.java b/...rc/main/java/software/amazon/smithy/aws/apigateway/openapi/AddCorsToGatewayResponses.java
@@ -109,22 +109,30 @@ private Node updateGatewayResponses(
         }
 
         return gatewayResponses.getMembers().entrySet().stream()
-                .collect(ObjectNode.collect(
-                        Map.Entry::getKey,
-                        entry -> updateGatewayResponse(corsHeaders, entry.getValue().expectObjectNode())));
+                .collect(ObjectNode.collect(Map.Entry::getKey, entry -> {
+                    return updateGatewayResponse(context, trait, corsHeaders, entry.getValue().expectObjectNode());
+                }));
     }
 
-    private ObjectNode updateGatewayResponse(Map<CorsHeader, String> sharedHeaders, ObjectNode gatewayResponse) {
+    private ObjectNode updateGatewayResponse(
+            Context<? extends Trait> context,
+            CorsTrait trait,
+            Map<CorsHeader, String> sharedHeaders,
+            ObjectNode gatewayResponse
+    ) {
         ObjectNode responseParameters = gatewayResponse
                 .getObjectMember(RESPONSE_PARAMETERS_KEY)
                 .orElse(Node.objectNode());
 
         // Track all CORS headers of the gateway response.
         Map<CorsHeader, String> headers = new TreeMap<>(sharedHeaders);
 
+        // Add the modeled additional headers. These could potentially be added by an
+        // apigateway feature, so they need to be present.
+        Set<String> exposedHeaders = new TreeSet<>(trait.getAdditionalExposedHeaders());
+
         // Find all headers exposed already in the response. These need to be added to the
         // Access-Control-Expose-Headers header if any are found.
-        Set<String> exposedHeaders = new TreeSet<>();
         for (String key : responseParameters.getStringMap().keySet()) {
             if (key.startsWith(HEADER_PREFIX)) {
                 exposedHeaders.add(key.substring(HEADER_PREFIX.length()));

diff --git a/.../src/test/resources/software/amazon/smithy/aws/apigateway/openapi/cors-model.openapi.json b/.../src/test/resources/software/amazon/smithy/aws/apigateway/openapi/cors-model.openapi.json
@@ -420,15 +420,17 @@
         "application/json": "{\"message\":$context.error.messageString}"
       },
       "responseParameters": {
-        "gatewayresponse.header.Access-Control-Allow-Origin": "'https://www.example.com'"
+        "gatewayresponse.header.Access-Control-Allow-Origin": "'https://www.example.com'",
+        "gatewayresponse.header.Access-Control-Expose-Headers": "'X-Service-Output-Metadata'"
       }
     },
     "DEFAULT_5XX": {
       "responseTemplates": {
         "application/json": "{\"message\":$context.error.messageString}"
       },
       "responseParameters": {
-        "gatewayresponse.header.Access-Control-Allow-Origin": "'https://www.example.com'"
+        "gatewayresponse.header.Access-Control-Allow-Origin": "'https://www.example.com'",
+        "gatewayresponse.header.Access-Control-Expose-Headers": "'X-Service-Output-Metadata'"
       }
     }
   }

diff --git a/...s/software/amazon/smithy/aws/apigateway/openapi/cors-with-additional-headers.openapi.json b/...s/software/amazon/smithy/aws/apigateway/openapi/cors-with-additional-headers.openapi.json
@@ -420,15 +420,17 @@
         "application/json": "{\"message\":$context.error.messageString}"
       },
       "responseParameters": {
-        "gatewayresponse.header.Access-Control-Allow-Origin": "'https://www.example.com'"
+        "gatewayresponse.header.Access-Control-Allow-Origin": "'https://www.example.com'",
+        "gatewayresponse.header.Access-Control-Expose-Headers": "'X-Service-Output-Metadata'"
       }
     },
     "DEFAULT_5XX": {
       "responseTemplates": {
         "application/json": "{\"message\":$context.error.messageString}"
       },
       "responseParameters": {
-        "gatewayresponse.header.Access-Control-Allow-Origin": "'https://www.example.com'"
+        "gatewayresponse.header.Access-Control-Allow-Origin": "'https://www.example.com'",
+        "gatewayresponse.header.Access-Control-Expose-Headers": "'X-Service-Output-Metadata'"
       }
     }
   }