feat: Expose sigv4 config for presigner (#402)

wooj2 · web-flow · commit 8f306bd62989 · 2021-09-21T14:18:36.000-07:00
diff --git a/AWSClientRuntime/Sources/Middlewares/SigV4Middleware.swift b/AWSClientRuntime/Sources/Middlewares/SigV4Middleware.swift
@@ -58,6 +58,7 @@ public struct SigV4Middleware<OperationStackOutput: HttpResponseBinding,
             let credentials = try credentialsProvider.getCredentials()
             
             let signingConfig = AWSSigningConfig(credentials: credentials,
+                                                 expiration: config.expiration,
                                                  signedBodyHeader: config.signedBodyHeader,
                                                  signedBodyValue: signedBodyValue,
                                                  flags: flags,
diff --git a/AWSClientRuntime/Sources/Signing/SigV4Config.swift b/AWSClientRuntime/Sources/Signing/SigV4Config.swift
@@ -12,6 +12,7 @@ public struct SigV4Config {
     let useDoubleURIEncode: Bool
     let shouldNormalizeURIPath: Bool
     let omitSessionToken: Bool
+    let expiration: Int64
     let signedBodyHeader: AWSSignedBodyHeader
     let unsignedBody: Bool
     
@@ -21,6 +22,7 @@ public struct SigV4Config {
                 useDoubleURIEncode: Bool = true,
                 shouldNormalizeURIPath: Bool = true,
                 omitSessionToken: Bool = false,
+                expiration: Int64 = 0,
                 signedBodyHeader: AWSSignedBodyHeader = .none,
                 unsignedBody: Bool) {
         self.credentialsProvider = credentialsProvider
@@ -29,6 +31,7 @@ public struct SigV4Config {
         self.useDoubleURIEncode = useDoubleURIEncode
         self.shouldNormalizeURIPath = shouldNormalizeURIPath
         self.omitSessionToken = omitSessionToken
+        self.expiration = expiration
         self.signedBodyHeader = signedBodyHeader
         self.unsignedBody = unsignedBody
     }
diff --git a/codegen/smithy-aws-swift-codegen/src/main/kotlin/software/amazon/smithy/aws/swift/codegen/PresignerGenerator.kt b/codegen/smithy-aws-swift-codegen/src/main/kotlin/software/amazon/smithy/aws/swift/codegen/PresignerGenerator.kt
@@ -1,9 +1,9 @@
 package software.amazon.smithy.aws.swift.codegen
 
 import software.amazon.smithy.aws.swift.codegen.AWSClientRuntimeTypes.Core.AWSClientConfiguration
+import software.amazon.smithy.aws.swift.codegen.AWSClientRuntimeTypes.Signing.SigV4Config
 import software.amazon.smithy.aws.swift.codegen.middleware.AWSSigningMiddleware
 import software.amazon.smithy.aws.swift.codegen.model.traits.Presignable
-import software.amazon.smithy.model.knowledge.OperationIndex
 import software.amazon.smithy.model.shapes.OperationShape
 import software.amazon.smithy.model.shapes.ServiceShape
 import software.amazon.smithy.swift.codegen.ClientRuntimeTypes.Http.SdkHttpRequest
@@ -14,6 +14,7 @@ import software.amazon.smithy.swift.codegen.core.CodegenContext
 import software.amazon.smithy.swift.codegen.core.toProtocolGenerationContext
 import software.amazon.smithy.swift.codegen.integration.SwiftIntegration
 import software.amazon.smithy.swift.codegen.middleware.MiddlewareExecutionGenerator
+import software.amazon.smithy.swift.codegen.middleware.MiddlewareRenderableExecutionContext
 import software.amazon.smithy.swift.codegen.model.expectShape
 
 data class PresignableOperation(
@@ -58,11 +59,10 @@ class PresignerGenerator : SwiftIntegration {
         writer.addImport(AWSClientConfiguration)
         writer.addImport(SdkHttpRequest)
 
-        val operationsIndex = OperationIndex.of(ctx.model)
         val httpBindingResolver = protocolGenerator.getProtocolHttpBindingResolver(protocolGeneratorContext, protocolGenerator.defaultContentType)
 
         writer.openBlock("extension $inputType {", "}") {
-            writer.openBlock("public func presign(config: \$N) -> \$T {", "}", AWSClientConfiguration, SdkHttpRequest) {
+            writer.openBlock("public func presign(config: \$N, sigv4Config: \$D) -> \$T {", "}", AWSClientConfiguration, SigV4Config, SdkHttpRequest) {
                 writer.write("let serviceName = \"${ctx.settings.sdkId}\"")
                 writer.write("let input = self")
                 val operationStackName = "operation"
@@ -78,7 +78,8 @@ class PresignerGenerator : SwiftIntegration {
                     httpBindingResolver,
                     protocolGenerator.httpProtocolCustomizable,
                     protocolGenerator.operationMiddleware,
-                    operationStackName
+                    operationStackName,
+                    MiddlewareRenderableExecutionContext.PRESIGNER
                 )
                 generator.render(op) { writer, _ ->
                     writer.write("return nil")
diff --git a/codegen/smithy-aws-swift-codegen/src/main/kotlin/software/amazon/smithy/aws/swift/codegen/middleware/AWSSigningMiddleware.kt b/codegen/smithy-aws-swift-codegen/src/main/kotlin/software/amazon/smithy/aws/swift/codegen/middleware/AWSSigningMiddleware.kt
@@ -18,6 +18,7 @@ import software.amazon.smithy.model.traits.OptionalAuthTrait
 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.MiddlewareRenderableExecutionContext
 import software.amazon.smithy.swift.codegen.middleware.MiddlewareStep
 import software.amazon.smithy.swift.codegen.model.expectTrait
 import software.amazon.smithy.swift.codegen.model.hasTrait
@@ -37,18 +38,22 @@ open class AWSSigningMiddleware(val paramsCallback: AWSSigningMiddlewareParamsCa
         symbolProvider: SymbolProvider,
         writer: SwiftWriter,
         op: OperationShape,
-        operationStackName: String
+        operationStackName: String,
+        executionContext: MiddlewareRenderableExecutionContext
     ) {
-        renderConfigDeclaration(writer, op)
+        renderConfigDeclaration(writer, op, executionContext)
         writer.write(
             "$operationStackName.${middlewareStep.stringValue()}.intercept(position: ${position.stringValue()}, middleware: \$N(config: sigv4Config))",
             AWSClientRuntimeTypes.Signing.SigV4Middleware
         )
     }
 
-    private fun renderConfigDeclaration(writer: SwiftWriter, op: OperationShape) {
+    private fun renderConfigDeclaration(writer: SwiftWriter, op: OperationShape, executionContext: MiddlewareRenderableExecutionContext) {
         writer.addImport(SigV4Config)
-        writer.write("let sigv4Config = \$N(${middlewareParamsString(op)})", SigV4Config)
+        when (executionContext) {
+            MiddlewareRenderableExecutionContext.PRESIGNER -> writer.write("let sigv4Config = sigv4Config ?? \$N(${middlewareParamsString(op)})", SigV4Config)
+            else -> writer.write("let sigv4Config = \$N(${middlewareParamsString(op)})", SigV4Config)
+        }
     }
 
     private fun middlewareParamsString(op: OperationShape): String {
diff --git a/codegen/smithy-aws-swift-codegen/src/main/kotlin/software/amazon/smithy/aws/swift/codegen/middleware/AWSXAmzTargetMiddleware.kt b/codegen/smithy-aws-swift-codegen/src/main/kotlin/software/amazon/smithy/aws/swift/codegen/middleware/AWSXAmzTargetMiddleware.kt
@@ -14,6 +14,7 @@ import software.amazon.smithy.swift.codegen.ServiceGenerator
 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.MiddlewareRenderableExecutionContext
 import software.amazon.smithy.swift.codegen.middleware.MiddlewareStep
 
 class AWSXAmzTargetMiddleware(val serviceShape: ServiceShape) : MiddlewareRenderable {
@@ -29,7 +30,8 @@ class AWSXAmzTargetMiddleware(val serviceShape: ServiceShape) : MiddlewareRender
         symbolProvider: SymbolProvider,
         writer: SwiftWriter,
         op: OperationShape,
-        operationStackName: String
+        operationStackName: String,
+        executionContext: MiddlewareRenderableExecutionContext
     ) {
         val inputShapeName = ServiceGenerator.getOperationInputShapeName(symbolProvider, model, op)
         val outputShapeName = ServiceGenerator.getOperationOutputShapeName(symbolProvider, model, op)
diff --git a/codegen/smithy-aws-swift-codegen/src/main/kotlin/software/amazon/smithy/aws/swift/codegen/middleware/EndpointResolverMiddleware.kt b/codegen/smithy-aws-swift-codegen/src/main/kotlin/software/amazon/smithy/aws/swift/codegen/middleware/EndpointResolverMiddleware.kt
@@ -12,6 +12,7 @@ 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.MiddlewareRenderableExecutionContext
 import software.amazon.smithy.swift.codegen.middleware.MiddlewareStep
 
 class EndpointResolverMiddleware : MiddlewareRenderable {
@@ -22,7 +23,7 @@ class EndpointResolverMiddleware : MiddlewareRenderable {
 
     override val position = MiddlewarePosition.BEFORE
 
-    override fun render(model: Model, symbolProvider: SymbolProvider, writer: SwiftWriter, op: OperationShape, operationStackName: String) {
+    override fun render(model: Model, symbolProvider: SymbolProvider, writer: SwiftWriter, op: OperationShape, operationStackName: String, executionContext: MiddlewareRenderableExecutionContext) {
         writer.write("$operationStackName.${middlewareStep.stringValue()}.intercept(position: ${position.stringValue()}, middleware: \$N(${middlewareParamsString()}))", AWSClientRuntimeTypes.Core.EndpointResolverMiddleware)
     }
 
diff --git a/codegen/smithy-aws-swift-codegen/src/main/kotlin/software/amazon/smithy/aws/swift/codegen/middleware/MutateHeadersMiddleware.kt b/codegen/smithy-aws-swift-codegen/src/main/kotlin/software/amazon/smithy/aws/swift/codegen/middleware/MutateHeadersMiddleware.kt
@@ -7,6 +7,7 @@ import software.amazon.smithy.swift.codegen.ClientRuntimeTypes
 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.MiddlewareRenderableExecutionContext
 import software.amazon.smithy.swift.codegen.middleware.MiddlewareStep
 
 class MutateHeadersMiddleware(
@@ -20,7 +21,7 @@ class MutateHeadersMiddleware(
 
     override val position = MiddlewarePosition.AFTER
 
-    override fun render(model: Model, symbolProvider: SymbolProvider, writer: SwiftWriter, op: OperationShape, operationStackName: String) {
+    override fun render(model: Model, symbolProvider: SymbolProvider, writer: SwiftWriter, op: OperationShape, operationStackName: String, executionContext: MiddlewareRenderableExecutionContext) {
         val paramsString = middlewareParamsString()
         writer.write("$operationStackName.${middlewareStep.stringValue()}.intercept(position: ${position.stringValue()}, middleware: \$N($paramsString))", ClientRuntimeTypes.Middleware.MutateHeadersMiddleware)
     }
diff --git a/codegen/smithy-aws-swift-codegen/src/main/kotlin/software/amazon/smithy/aws/swift/codegen/middleware/RetryMiddleware.kt b/codegen/smithy-aws-swift-codegen/src/main/kotlin/software/amazon/smithy/aws/swift/codegen/middleware/RetryMiddleware.kt
@@ -12,6 +12,7 @@ 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.MiddlewareRenderableExecutionContext
 import software.amazon.smithy.swift.codegen.middleware.MiddlewareStep
 
 class RetryMiddleware : MiddlewareRenderable {
@@ -22,7 +23,7 @@ class RetryMiddleware : MiddlewareRenderable {
 
     override val position = MiddlewarePosition.AFTER
 
-    override fun render(model: Model, symbolProvider: SymbolProvider, writer: SwiftWriter, op: OperationShape, operationStackName: String) {
+    override fun render(model: Model, symbolProvider: SymbolProvider, writer: SwiftWriter, op: OperationShape, operationStackName: String, executionContext: MiddlewareRenderableExecutionContext) {
         writer.write("$operationStackName.${middlewareStep.stringValue()}.intercept(position: ${position.stringValue()}, middleware: \$N(${middlewareParamsString()}))", AWSClientRuntimeTypes.Core.RetryerMiddleware)
     }
 
diff --git a/codegen/smithy-aws-swift-codegen/src/main/kotlin/software/amazon/smithy/aws/swift/codegen/middleware/UserAgentMiddleware.kt b/codegen/smithy-aws-swift-codegen/src/main/kotlin/software/amazon/smithy/aws/swift/codegen/middleware/UserAgentMiddleware.kt
@@ -13,6 +13,7 @@ import software.amazon.smithy.swift.codegen.SwiftSettings
 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.MiddlewareRenderableExecutionContext
 import software.amazon.smithy.swift.codegen.middleware.MiddlewareStep
 
 class UserAgentMiddleware(val settings: SwiftSettings) : MiddlewareRenderable {
@@ -23,7 +24,7 @@ class UserAgentMiddleware(val settings: SwiftSettings) : MiddlewareRenderable {
 
     override val position = MiddlewarePosition.BEFORE
 
-    override fun render(model: Model, symbolProvider: SymbolProvider, writer: SwiftWriter, op: OperationShape, operationStackName: String) {
+    override fun render(model: Model, symbolProvider: SymbolProvider, writer: SwiftWriter, op: OperationShape, operationStackName: String, executionContext: MiddlewareRenderableExecutionContext) {
         writer.write("$operationStackName.${middlewareStep.stringValue()}.intercept(position: ${position.stringValue()}, middleware: \$N(${middlewareParamsString()}))", AWSClientRuntimeTypes.Core.UserAgentMiddleware)
     }
 
diff --git a/codegen/smithy-aws-swift-codegen/src/test/kotlin/software/amazon/smithy/aws/swift/codegen/AWSSigningMiddlewareTests.kt b/codegen/smithy-aws-swift-codegen/src/test/kotlin/software/amazon/smithy/aws/swift/codegen/AWSSigningMiddlewareTests.kt
@@ -20,6 +20,7 @@ import software.amazon.smithy.model.traits.AuthTrait
 import software.amazon.smithy.model.traits.HttpBasicAuthTrait
 import software.amazon.smithy.model.traits.OptionalAuthTrait
 import software.amazon.smithy.swift.codegen.SwiftWriter
+import software.amazon.smithy.swift.codegen.middleware.MiddlewareRenderableExecutionContext
 
 class AWSSigningMiddlewareTests {
     @Test
@@ -87,7 +88,7 @@ stack.finalizeStep.intercept(position: .before, middleware: AWSClientRuntime.Sig
         val opStackName = "stack"
         val sut = AWSSigningMiddleware()
 
-        sut.render(context.model, context.symbolProvider, writer, operationShape, opStackName)
+        sut.render(context.model, context.symbolProvider, writer, operationShape, opStackName, MiddlewareRenderableExecutionContext.CLIENT)
 
         val contents = writer.toString()
         contents.shouldContainOnlyOnce(expectedContents)
@@ -113,7 +114,60 @@ stack.finalizeStep.intercept(position: .before, middleware: AWSClientRuntime.Sig
         val opStackName = "stack"
         val sut = AWSSigningMiddleware()
 
-        sut.render(context.model, context.symbolProvider, writer, operationShape, opStackName)
+        sut.render(context.model, context.symbolProvider, writer, operationShape, opStackName, MiddlewareRenderableExecutionContext.CLIENT)
+
+        val contents = writer.toString()
+        contents.shouldContainOnlyOnce(expectedContents)
+    }
+
+    @Test
+    fun `renderSigningMiddleware unsignedBody true, presigner`() {
+        val expectedContents =
+            """
+let sigv4Config = sigv4Config ?? AWSClientRuntime.SigV4Config(unsignedBody: true)
+stack.finalizeStep.intercept(position: .before, middleware: AWSClientRuntime.SigV4Middleware(config: sigv4Config))"""
+        val writer = SwiftWriter("testName")
+        val serviceShape = ServiceShape.builder()
+            .id("com.test#Example")
+            .version("1.0")
+            .addTrait(SigV4Trait.builder().name("ExampleService").build())
+            .build()
+        val operationShape = OperationShape.builder()
+            .id("com.test#ExampleOperation")
+            .addTrait(UnsignedPayloadTrait())
+            .build()
+        val model = Model.builder().addShape(serviceShape).addShape(operationShape).build()
+        val context = model.newTestContext(generator = AWSRestJson1ProtocolGenerator()).ctx
+        val opStackName = "stack"
+        val sut = AWSSigningMiddleware()
+
+        sut.render(context.model, context.symbolProvider, writer, operationShape, opStackName, MiddlewareRenderableExecutionContext.PRESIGNER)
+
+        val contents = writer.toString()
+        contents.shouldContainOnlyOnce(expectedContents)
+    }
+
+    @Test
+    fun `renderSigningMiddleware unsignedBody false, presigner`() {
+        val expectedContents =
+            """
+let sigv4Config = sigv4Config ?? AWSClientRuntime.SigV4Config(unsignedBody: false)
+stack.finalizeStep.intercept(position: .before, middleware: AWSClientRuntime.SigV4Middleware(config: sigv4Config))"""
+        val writer = SwiftWriter("testName")
+        val serviceShape = ServiceShape.builder()
+            .id("com.test#Example")
+            .version("1.0")
+            .addTrait(SigV4Trait.builder().name("ExampleService").build())
+            .build()
+        val operationShape = OperationShape.builder()
+            .id("com.test#ExampleOperation")
+            .build()
+        val model = Model.builder().addShape(serviceShape).addShape(operationShape).build()
+        val context = model.newTestContext(generator = AWSRestJson1ProtocolGenerator()).ctx
+        val opStackName = "stack"
+        val sut = AWSSigningMiddleware()
+
+        sut.render(context.model, context.symbolProvider, writer, operationShape, opStackName, MiddlewareRenderableExecutionContext.PRESIGNER)
 
         val contents = writer.toString()
         contents.shouldContainOnlyOnce(expectedContents)
diff --git a/codegen/smithy-aws-swift-codegen/src/test/kotlin/software/amazon/smithy/aws/swift/codegen/AWSXAmzTargetMiddlewareTests.kt b/codegen/smithy-aws-swift-codegen/src/test/kotlin/software/amazon/smithy/aws/swift/codegen/AWSXAmzTargetMiddlewareTests.kt
@@ -22,6 +22,7 @@ import software.amazon.smithy.swift.codegen.SwiftDelegator
 import software.amazon.smithy.swift.codegen.SwiftSettings
 import software.amazon.smithy.swift.codegen.SwiftWriter
 import software.amazon.smithy.swift.codegen.integration.ProtocolGenerator
+import software.amazon.smithy.swift.codegen.middleware.MiddlewareRenderableExecutionContext
 
 class AWSXAmzTargetMiddlewareTests {
     @Test
@@ -56,7 +57,7 @@ stack.serializeStep.intercept(position: .before, middleware: AWSClientRuntime.XA
         val context = model.newTestContext("com.test#ExampleServiceShapeName", AwsJson1_0_ProtocolGenerator()).ctx
         val sut = AWSXAmzTargetMiddleware(context.service)
 
-        sut.render(context.model, context.symbolProvider, writer, operationShape, opStackName)
+        sut.render(context.model, context.symbolProvider, writer, operationShape, opStackName, MiddlewareRenderableExecutionContext.CLIENT)
 
         val contents = writer.toString()
         contents.shouldContainOnlyOnce(expectedContents)
diff --git a/codegen/smithy-aws-swift-codegen/src/test/kotlin/software/amazon/smithy/aws/swift/codegen/PresignerGeneratorTests.kt b/codegen/smithy-aws-swift-codegen/src/test/kotlin/software/amazon/smithy/aws/swift/codegen/PresignerGeneratorTests.kt
