Fix Access-Control-Expose-Headers

smithy-aws-apigateway-openapi/src/main/java/software/amazon/smithy/aws/apigateway/openapi/AddCorsResponseHeaders.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License").
  * You may not use this file except in compliance with the License.
@@ -17,12 +17,14 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.logging.Logger;
 import software.amazon.smithy.jsonschema.Schema;
 import software.amazon.smithy.model.shapes.OperationShape;
 import software.amazon.smithy.model.traits.CorsTrait;
 import software.amazon.smithy.model.traits.Trait;
 import software.amazon.smithy.openapi.fromsmithy.Context;
+import software.amazon.smithy.openapi.model.OperationObject;
 import software.amazon.smithy.openapi.model.ParameterObject;
 import software.amazon.smithy.openapi.model.Ref;
 import software.amazon.smithy.openapi.model.ResponseObject;
@@ -48,25 +50,43 @@ public List<ApiGatewayConfig.ApiType> getApiTypes() {
     }
 
     @Override
-    public ResponseObject updateResponse(
+    public OperationObject postProcessOperation(
             Context<? extends Trait> context,
             OperationShape shape,
-            String status,
+            OperationObject operation,
             String method,
-            String path,
-            ResponseObject response
+            String path
     ) {
         return context.getService().getTrait(CorsTrait.class)
-                .map(corsTrait -> addCorsHeadersToResponse(context, shape, response, corsTrait, method))
-                .orElse(response);
+                .map(trait -> addCorsHeadersToResponses(context, shape, operation, method, trait))
+                .orElse(operation);
     }
 
-    private ResponseObject addCorsHeadersToResponse(
+    private OperationObject addCorsHeadersToResponses(
+            Context<? extends Trait> context,
+            OperationShape shape,
+            OperationObject operationObject,
+            String method,
+            CorsTrait trait
+    ) {
+        OperationObject.Builder builder = operationObject.toBuilder();
+
+        for (Map.Entry<String, ResponseObject> entry : operationObject.getResponses().entrySet()) {
+            ResponseObject updated = createUpdatedResponseWithCorsHeaders(
+                    context, shape, operationObject, method, trait, entry.getValue());
+            builder.putResponse(entry.getKey(), updated);
+        }
+
+        return builder.build();
+    }
+
+    private ResponseObject createUpdatedResponseWithCorsHeaders(
             Context<? extends Trait> context,
             OperationShape operation,
-            ResponseObject response,
-            CorsTrait corsTrait,
-            String method
+            OperationObject operationObject,
+            String method,
+            CorsTrait trait,
+            ResponseObject response
     ) {
         // Determine which headers have been added to the response.
         List<String> headers = new ArrayList<>();
@@ -108,7 +128,7 @@ private ResponseObject addCorsHeadersToResponse(
             // An HTTP response to a CORS request that is *not* a
             // CORS-preflight request can also include the following header:
             // `Access-Control-Expose-Headers`
-            if (!CorsHeader.deduceOperationHeaders(context, operation, corsTrait).isEmpty()) {
+            if (!CorsHeader.deduceOperationResponseHeaders(context, operationObject, operation, trait).isEmpty()) {
                 headers.add(CorsHeader.EXPOSE_HEADERS.toString());
             }
 

smithy-aws-apigateway-openapi/src/main/java/software/amazon/smithy/aws/apigateway/openapi/AddIntegrations.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License").
  * You may not use this file except in compliance with the License.
@@ -74,7 +74,7 @@ public OperationObject updateOperation(
         IntegrationTraitIndex index = IntegrationTraitIndex.of(context.getModel());
         return index.getIntegrationTrait(context.getService(), shape)
                 .map(trait -> operation.toBuilder()
-                        .putExtension(EXTENSION_NAME, createIntegration(context, shape, trait))
+                        .putExtension(EXTENSION_NAME, createIntegration(context, operation, shape, trait))
                         .build())
                 .orElseGet(() -> {
                     LOGGER.warning("No API Gateway integration trait found for " + shape.getId());
@@ -84,14 +84,15 @@ public OperationObject updateOperation(
 
     private ObjectNode createIntegration(
             Context<? extends Trait> context,
+            OperationObject operationObject,
             OperationShape shape,
             Trait integration
     ) {
         ObjectNode integrationObject = getIntegrationAsObject(context, shape, integration);
         return context.getService().getTrait(CorsTrait.class)
                 .map(cors -> {
                     LOGGER.fine(() -> String.format("Adding CORS to `%s` operation responses", shape.getId()));
-                    return updateIntegrationWithCors(context, shape, integrationObject, cors);
+                    return updateIntegrationWithCors(context, operationObject, shape, integrationObject, cors);
                 })
                 .orElse(integrationObject);
     }
@@ -122,6 +123,7 @@ private static ObjectNode getIntegrationAsObject(
 
     private ObjectNode updateIntegrationWithCors(
             Context<? extends Trait> context,
+            OperationObject operationObject,
             OperationShape shape,
             ObjectNode integrationNode,
             CorsTrait cors
@@ -141,7 +143,7 @@ private ObjectNode updateIntegrationWithCors(
 
         LOGGER.finer(() -> String.format("Adding the following CORS headers to the API Gateway integration of %s: %s",
                                          shape.getId(), corsHeaders));
-        Set<String> deducedHeaders = CorsHeader.deduceOperationHeaders(context, shape, cors);
+        Set<String> deducedHeaders = CorsHeader.deduceOperationResponseHeaders(context, operationObject, shape, cors);
         LOGGER.fine(() -> String.format("Detected the following headers for operation %s: %s",
                                         shape.getId(), deducedHeaders));
 

smithy-aws-apigateway-openapi/src/main/java/software/amazon/smithy/aws/apigateway/openapi/ApiGatewayMapper.java

@@ -112,6 +112,18 @@ public OperationObject updateOperation(
                        : operation;
             }
 
+            @Override
+            public OperationObject postProcessOperation(
+                    Context<? extends Trait> context,
+                    OperationShape shape,
+                    OperationObject operation,
+                    String httpMethodName, String path
+            ) {
+                return matchesApiType(context)
+                       ? delegate.postProcessOperation(context, shape, operation, httpMethodName, path)
+                       : operation;
+            }
+
             @Override
             public PathItem updatePathItem(Context<? extends Trait> context, String path, PathItem pathItem) {
                 return matchesApiType(context)

smithy-aws-apigateway-openapi/src/main/java/software/amazon/smithy/aws/apigateway/openapi/CorsHeader.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License").
  * You may not use this file except in compliance with the License.
@@ -22,6 +22,8 @@
 import software.amazon.smithy.model.traits.Trait;
 import software.amazon.smithy.openapi.fromsmithy.Context;
 import software.amazon.smithy.openapi.fromsmithy.SecuritySchemeConverter;
+import software.amazon.smithy.openapi.model.OperationObject;
+import software.amazon.smithy.openapi.model.ResponseObject;
 
 enum CorsHeader {
 
@@ -43,8 +45,9 @@ public String toString() {
         return headerName;
     }
 
-    static <T extends Trait> Set<String> deduceOperationHeaders(
+    static <T extends Trait> Set<String> deduceOperationResponseHeaders(
             Context<T> context,
+            OperationObject operationObject,
             OperationShape shape,
             CorsTrait cors
     ) {
@@ -58,6 +61,11 @@ static <T extends Trait> Set<String> deduceOperationHeaders(
             result.addAll(getSecuritySchemeResponseHeaders(context, converter));
         }
 
+        // Include all headers found in the generated OpenAPI response.
+        for (ResponseObject responseObject : operationObject.getResponses().values()) {
+            result.addAll(responseObject.getHeaders().keySet());
+        }
+
         return result;
     }
 

smithy-aws-apigateway-openapi/src/test/resources/software/amazon/smithy/aws/apigateway/openapi/cors-explicit-options.openapi.json

@@ -63,7 +63,8 @@
             "default": {
               "statusCode": "200",
               "responseParameters": {
-                "method.response.header.Access-Control-Allow-Origin": "'https://foo.com'"
+                "method.response.header.Access-Control-Allow-Origin": "'https://foo.com'",
+                "method.response.header.Access-Control-Expose-Headers": "'X-Hd'"
               }
             }
           }

smithy-aws-apigateway-openapi/src/test/resources/software/amazon/smithy/aws/apigateway/openapi/cors-model.openapi.json

@@ -199,7 +199,7 @@
             "default": {
               "statusCode": "200",
               "responseParameters": {
-                "method.response.header.Access-Control-Expose-Headers": "'X-Service-Output-Metadata'",
+                "method.response.header.Access-Control-Expose-Headers": "'X-Foo-Header,X-Service-Output-Metadata'",
                 "method.response.header.Access-Control-Allow-Origin": "'https://www.example.com'"
               }
             }

smithy-aws-apigateway-openapi/src/test/resources/software/amazon/smithy/aws/apigateway/openapi/cors-with-additional-headers.openapi.json

@@ -199,7 +199,7 @@
             "default": {
               "statusCode": "200",
               "responseParameters": {
-                "method.response.header.Access-Control-Expose-Headers": "'X-Service-Output-Metadata'",
+                "method.response.header.Access-Control-Expose-Headers": "'X-Foo-Header,X-Service-Output-Metadata'",
                 "method.response.header.Access-Control-Allow-Origin": "'https://www.example.com'"
               }
             }

smithy-aws-apigateway-openapi/src/test/resources/software/amazon/smithy/aws/apigateway/openapi/cors-with-custom-gateway-response-headers.openapi.json

@@ -63,7 +63,8 @@
             "default": {
               "statusCode": "200",
               "responseParameters": {
-                "method.response.header.Access-Control-Allow-Origin": "'https://foo.com'"
+                "method.response.header.Access-Control-Allow-Origin": "'https://foo.com'",
+                "method.response.header.Access-Control-Expose-Headers": "'X-Hd'"
               }
             }
           }

smithy-openapi/src/main/java/software/amazon/smithy/openapi/fromsmithy/OpenApiConverter.java

@@ -426,10 +426,12 @@ private <T extends Trait> void addPaths(
                 String method = result.getMethod();
                 String path = result.getUri();
                 PathItem.Builder pathItem = paths.computeIfAbsent(result.getUri(), (uri) -> PathItem.builder());
+
                 // Mark the operation deprecated if the trait's present.
                 if (shape.hasTrait(DeprecatedTrait.class)) {
                     result.getOperation().deprecated(true);
                 }
+
                 // Add security requirements to the operation.
                 addOperationSecurity(context, result.getOperation(), shape, plugin);
 
@@ -438,9 +440,10 @@ private <T extends Trait> void addPaths(
                         .map(DocumentationTrait::getValue)
                         .ifPresent(description -> result.getOperation().description(description));
 
-                // Pass the operation through the plugin system and then build it.
-                OperationObject builtOperation = plugin.updateOperation(
-                        context, shape, result.getOperation().build(), method, path);
+                OperationObject builtOperation = result.getOperation().build();
+
+                // Pass the operation through the plugin system.
+                builtOperation = plugin.updateOperation(context, shape, builtOperation, method, path);
                 // Add tags that are on the operation.
                 builtOperation = addOperationTags(context, shape, builtOperation);
                 // Update each parameter of the operation and rebuild if necessary.
@@ -449,6 +452,8 @@ private <T extends Trait> void addPaths(
                 builtOperation = updateResponses(context, shape, builtOperation, method, path, plugin);
                 // Update the request body of the operation and rebuild if necessary.
                 builtOperation = updateRequestBody(context, shape, builtOperation, method, path, plugin);
+                // Pass the operation through the plugin system for post-processing.
+                builtOperation = plugin.postProcessOperation(context, shape, builtOperation, method, path);
 
                 switch (method.toLowerCase(Locale.ENGLISH)) {
                     case "get":

smithy-openapi/src/main/java/software/amazon/smithy/openapi/fromsmithy/OpenApiMapper.java

@@ -71,7 +71,10 @@ default byte getOrder() {
     default void updateDefaultSettings(Model model, OpenApiConfig config) {}
 
     /**
-     * Updates an operation.
+     * Updates an operation before invoking the plugin system on the contents
+     * of the operation (specifically, before {@link #updateParameter},
+     * {@link #updateRequestBody}, {@link #updateResponse},
+     * {@link #updateRequestBody}, and {@link #postProcessOperation}).
      *
      * @param context Conversion context.
      * @param shape Operation being converted.
@@ -90,6 +93,29 @@ default OperationObject updateOperation(
         return operation;
     }
 
+    /**
+     * Updates an operation after invoking the plugin system on the contents
+     * of the operation (specifically, after {@link #updateOperation},
+     * {@link #updateParameter}, {@link #updateRequestBody},
+     * {@link #updateResponse}, and {@link #updateRequestBody}).
+     *
+     * @param context Conversion context.
+     * @param shape Operation being converted.
+     * @param operation OperationObject being built.
+     * @param httpMethodName The HTTP method of the operation.
+     * @param path The HTTP URI of the operation.
+     * @return Returns the updated operation object.
+     */
+    default OperationObject postProcessOperation(
+            Context<? extends Trait> context,
+            OperationShape shape,
+            OperationObject operation,
+            String httpMethodName,
+            String path
+    ) {
+        return operation;
+    }
+
     /**
      * Updates a path item.
      *
@@ -269,6 +295,23 @@ public OperationObject updateOperation(
                 return operation;
             }
 
+            @Override
+            public OperationObject postProcessOperation(
+                    Context<? extends Trait> context,
+                    OperationShape shape,
+                    OperationObject operation,
+                    String httpMethodName,
+                    String path
+            ) {
+                for (OpenApiMapper plugin : sorted) {
+                    if (operation == null) {
+                        return null;
+                    }
+                    operation = plugin.postProcessOperation(context, shape, operation, httpMethodName, path);
+                }
+                return operation;
+            }
+
             @Override
             public PathItem updatePathItem(Context<? extends Trait> context, String path, PathItem pathItem) {
                 for (OpenApiMapper plugin : sorted) {