feat: Customize the S3 putObject presigned URL to include S3 metada…

…ta (#1139)

IntegrationTests/Services/AWSS3IntegrationTests/S3URLEncodingTests.swift

@@ -28,18 +28,21 @@ final class S3URLEncodingTests: S3XCTestCase {
         }
     }
 
-    func test_presignPutObject_putsAllKeys() async throws {
+    func test_presignedURL_putObject_putsAllKeysWithMetadata() async throws {
         let config = try await S3Client.S3ClientConfiguration(region: region)
         for key in keys {
-            let input = PutObjectInput(body: .data(Data()), bucket: bucketName, key: key)
+            let input = PutObjectInput(body: .data(Data()), bucket: bucketName, key: key, metadata: ["filename": key])
             let presignedURLOrNil = try await input.presignURL(config: config, expiration: 30.0)
             let presignedURL = try XCTUnwrap(presignedURLOrNil)
             var urlRequest = URLRequest(url: presignedURL)
             urlRequest.httpMethod = "PUT"
             urlRequest.httpBody = Data()
             _ = try await perform(urlRequest: urlRequest)
         }
-        let createdKeys = Set(try await listBucketKeys())
-        XCTAssertTrue(createdKeys.isSuperset(of: keys))
+        for key in keys {
+            let input = HeadObjectInput(bucket: bucketName, key: key)
+            let output = try await client.headObject(input: input)
+            XCTAssertEqual(output.metadata?["filename"], key)
+        }
     }
 }

Package.swift

@@ -115,6 +115,8 @@ func addIntegrationTestTarget(_ name: String) {
             "README.md",
             "Resources/ECSIntegTestApp/"
         ]
+    case "AWSS3":
+        additionalDependencies = ["AWSSSOAdmin"]
     default:
         break
     }
@@ -552,4 +554,4 @@ let servicesWithIntegrationTests: [String] = [
 servicesWithIntegrationTests.forEach(addIntegrationTestTarget)
 
 // Uncomment this line to enable protocol tests
-// addProtocolTests()
+// addProtocolTests()

codegen/smithy-aws-swift-codegen/src/main/kotlin/software/amazon/smithy/aws/swift/codegen/customization/PutObjectPresignedURLMiddleware.kt

@@ -0,0 +1,36 @@
+package software.amazon.smithy.aws.swift.codegen.customization
+
+import software.amazon.smithy.codegen.core.Symbol
+import software.amazon.smithy.swift.codegen.Middleware
+import software.amazon.smithy.swift.codegen.SwiftWriter
+import software.amazon.smithy.swift.codegen.integration.steps.OperationSerializeStep
+
+// This middleware is intended only for use with S3 `PutObject`, and only for use when
+// creating a presigned URL.
+//
+// Generates a middleware that writes S3 object metadata into the HTTP query string.
+class PutObjectPresignedURLMiddleware(
+    inputSymbol: Symbol,
+    outputSymbol: Symbol,
+    outputErrorSymbol: Symbol,
+    private val writer: SwiftWriter
+) : Middleware(writer, inputSymbol, OperationSerializeStep(inputSymbol, outputSymbol, outputErrorSymbol)) {
+    override val typeName = "PutObjectPresignedURLMiddleware"
+
+    override fun generateInit() {
+        writer.write("public init() {}")
+    }
+
+    override fun generateMiddlewareClosure() {
+        writer.apply {
+            write("let metadata = input.operationInput.metadata ?? [:]")
+            openBlock("for (metadataKey, metadataValue) in metadata {", "}") {
+                openBlock("let queryItem = URLQueryItem(", ")") {
+                    write("name: \"x-amz-meta-\\(metadataKey.urlPercentEncoding())\",")
+                    write("value: metadataValue.urlPercentEncoding()")
+                }
+                write("input.builder.withQueryItem(queryItem)")
+            }
+        }
+    }
+}

codegen/smithy-aws-swift-codegen/src/main/kotlin/software/amazon/smithy/aws/swift/codegen/customization/presignable/PresignableUrlIntegration.kt

@@ -6,8 +6,10 @@ import software.amazon.smithy.aws.swift.codegen.AWSSigningParams
 import software.amazon.smithy.aws.swift.codegen.PresignableOperation
 import software.amazon.smithy.aws.swift.codegen.SigningAlgorithm
 import software.amazon.smithy.aws.swift.codegen.customization.InputTypeGETQueryItemMiddleware
+import software.amazon.smithy.aws.swift.codegen.customization.PutObjectPresignedURLMiddleware
 import software.amazon.smithy.aws.swift.codegen.middleware.AWSSigningMiddleware
 import software.amazon.smithy.aws.swift.codegen.middleware.InputTypeGETQueryItemMiddlewareRenderable
+import software.amazon.smithy.aws.swift.codegen.middleware.PutObjectPresignedURLMiddlewareRenderable
 import software.amazon.smithy.codegen.core.Symbol
 import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.knowledge.OperationIndex
@@ -68,8 +70,13 @@ class PresignableUrlIntegration(private val presignedOperations: Map<String, Set
                 val serviceConfig = AWSServiceConfig(writer, protocolGenerationContext)
                 renderPresigner(writer, ctx, delegator, op, inputType, serviceConfig)
             }
-            if (presignableOperation.operationId != "com.amazonaws.s3#PutObject") {
-                renderMiddlewareClassForQueryString(ctx, delegator, op)
+            when (presignableOperation.operationId) {
+                "com.amazonaws.s3#GetObject", "com.amazonaws.polly#SynthesizeSpeech" -> {
+                    renderMiddlewareClassForQueryString(ctx, delegator, op)
+                }
+                "com.amazonaws.s3#PutObject" -> {
+                    renderMiddlewareClassForPutObject(ctx, delegator, op)
+                }
             }
         }
     }
@@ -160,10 +167,15 @@ class PresignableUrlIntegration(private val presignedOperations: Map<String, Set
             )
             operationMiddlewareCopy.appendMiddleware(op, AWSSigningMiddleware(context.model, context.symbolProvider, params))
         }
-
-        if (op.id.toString() != "com.amazonaws.s3#PutObject") {
-            operationMiddlewareCopy.removeMiddleware(op, MiddlewareStep.SERIALIZESTEP, "OperationInputBodyMiddleware")
-            operationMiddlewareCopy.appendMiddleware(op, InputTypeGETQueryItemMiddlewareRenderable(inputSymbol))
+        when (op.id.toString()) {
+            "com.amazonaws.s3#GetObject", "com.amazonaws.polly#SynthesizeSpeech" -> {
+                operationMiddlewareCopy.removeMiddleware(op, MiddlewareStep.SERIALIZESTEP, "OperationInputBodyMiddleware")
+                operationMiddlewareCopy.appendMiddleware(op, InputTypeGETQueryItemMiddlewareRenderable(inputSymbol))
+            }
+            "com.amazonaws.s3#PutObject" -> {
+                operationMiddlewareCopy.removeMiddleware(op, MiddlewareStep.SERIALIZESTEP, "OperationInputBodyMiddleware")
+                operationMiddlewareCopy.appendMiddleware(op, PutObjectPresignedURLMiddlewareRenderable())
+            }
         }
 
         return operationMiddlewareCopy
@@ -201,6 +213,36 @@ class PresignableUrlIntegration(private val presignedOperations: Map<String, Set
         }
     }
 
+    private fun renderMiddlewareClassForPutObject(codegenContext: CodegenContext, delegator: SwiftDelegator, op: OperationShape) {
+
+        val serviceShape = codegenContext.model.expectShape<ServiceShape>(codegenContext.settings.service)
+        val ctx = codegenContext.toProtocolGenerationContext(serviceShape, delegator)?.let { it } ?: run { return }
+
+        val opIndex = OperationIndex.of(ctx.model)
+        val inputShape = opIndex.getInput(op).get()
+        val outputShape = opIndex.getOutput(op).get()
+        val operationErrorName = MiddlewareShapeUtils.outputErrorSymbolName(op)
+        val inputSymbol = ctx.symbolProvider.toSymbol(inputShape)
+        val outputSymbol = ctx.symbolProvider.toSymbol(outputShape)
+        val outputErrorSymbol = Symbol.builder().name(operationErrorName).build()
+
+        val rootNamespace = ctx.settings.moduleName
+        val headerMiddlewareSymbol = Symbol.builder()
+            .definitionFile("./$rootNamespace/models/${inputSymbol.name}+QueryItemMiddlewareForPresignUrl.swift")
+            .name(inputSymbol.name)
+            .build()
+        delegator.useShapeWriter(headerMiddlewareSymbol) { writer ->
+            writer.addImport(SwiftDependency.CLIENT_RUNTIME.target)
+            val queryItemMiddleware = PutObjectPresignedURLMiddleware(
+                inputSymbol,
+                outputSymbol,
+                outputErrorSymbol,
+                writer
+            )
+            MiddlewareGenerator(writer, queryItemMiddleware).generate()
+        }
+    }
+
     private fun overrideHttpMethod(operation: OperationShape): String {
         return when (operation.id.toString()) {
             "com.amazonaws.s3#PutObject" -> "put"

codegen/smithy-aws-swift-codegen/src/main/kotlin/software/amazon/smithy/aws/swift/codegen/middleware/PutObjectPresignedURLMiddlewareRenderable.kt

@@ -0,0 +1,29 @@
+package software.amazon.smithy.aws.swift.codegen.middleware
+
+import software.amazon.smithy.model.shapes.OperationShape
+import software.amazon.smithy.swift.codegen.SwiftWriter
+import software.amazon.smithy.swift.codegen.middleware.MiddlewarePosition
+import software.amazon.smithy.swift.codegen.middleware.MiddlewareRenderable
+import software.amazon.smithy.swift.codegen.middleware.MiddlewareStep
+
+// This middleware renderer inserts a custom middleware named `PutObjectPresignedURLMiddleware`
+// into the operation stack.  It is only intended for use with S3 `PutObject` and only when
+// generating a pre-signed URL.
+class PutObjectPresignedURLMiddlewareRenderable : MiddlewareRenderable {
+
+    override val name = "PutObjectPresignedURLMiddleware"
+
+    override val middlewareStep = MiddlewareStep.SERIALIZESTEP
+
+    override val position = MiddlewarePosition.AFTER
+
+    override fun render(writer: SwiftWriter, op: OperationShape, operationStackName: String) {
+        writer.write(
+            "\$L.\$L.intercept(position: \$L, middleware: \$L())",
+            operationStackName,
+            middlewareStep.stringValue(),
+            position.stringValue(),
+            name
+        )
+    }
+}

codegen/smithy-aws-swift-codegen/src/test/kotlin/software/amazon/smithy/aws/swift/codegen/customizations/PresignableUrlIntegrationTests.kt

@@ -50,6 +50,55 @@ class PresignableUrlIntegrationTests {
         contents.shouldContainOnlyOnce(expectedContents)
     }
 
+    @Test
+    fun `S3 PutObject operation stack contains the PutObjectPresignedURLMiddleware`() {
+        val context = setupTests("presign-urls-s3.smithy", "com.amazonaws.s3#AmazonS3")
+        val contents = TestContextGenerator.getFileContents(context.manifest, "/Example/models/PutObjectInput+Presigner.swift")
+        contents.shouldSyntacticSanityCheck()
+        val expectedContents = """
+        operation.serializeStep.intercept(position: .after, middleware: PutObjectPresignedURLMiddleware())
+        """
+        contents.shouldContainOnlyOnce(expectedContents)
+    }
+
+    @Test
+    fun `S3 PutObject's PutObjectPresignedURLMiddleware is rendered`() {
+        val context = setupTests("presign-urls-s3.smithy", "com.amazonaws.s3#AmazonS3")
+        val contents = TestContextGenerator.getFileContents(context.manifest, "/Example/models/PutObjectInput+QueryItemMiddlewareForPresignUrl.swift")
+        contents.shouldSyntacticSanityCheck()
+        val expectedContents = """
+public struct PutObjectPresignedURLMiddleware: ClientRuntime.Middleware {
+    public let id: Swift.String = "PutObjectPresignedURLMiddleware"
+
+    public init() {}
+
+    public func handle<H>(context: Context,
+                  input: ClientRuntime.SerializeStepInput<PutObjectInput>,
+                  next: H) async throws -> ClientRuntime.OperationOutput<PutObjectOutputResponse>
+    where H: Handler,
+    Self.MInput == H.Input,
+    Self.MOutput == H.Output,
+    Self.Context == H.Context
+    {
+        let metadata = input.operationInput.metadata ?? [:]
+        for (metadataKey, metadataValue) in metadata {
+            let queryItem = URLQueryItem(
+                name: "x-amz-meta-\(metadataKey.urlPercentEncoding())",
+                value: metadataValue.urlPercentEncoding()
+            )
+            input.builder.withQueryItem(queryItem)
+        }
+        return try await next.handle(context: context, input: input)
+    }
+
+    public typealias MInput = ClientRuntime.SerializeStepInput<PutObjectInput>
+    public typealias MOutput = ClientRuntime.OperationOutput<PutObjectOutputResponse>
+    public typealias Context = ClientRuntime.HttpContext
+}
+"""
+        contents.shouldContainOnlyOnce(expectedContents)
+    }
+
     private fun setupTests(smithyFile: String, serviceShapeId: String): TestContext {
         val context = TestContextGenerator.initContextFrom(smithyFile, serviceShapeId, RestXmlTrait.ID)
         val presigner = PresignableUrlIntegration()