Service sdk constraints (#1338)

Co-authored-by: luigi <[email protected]>

Author: luigi617

codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/service/KtorStubGenerator.kt

@@ -10,6 +10,8 @@ import software.amazon.smithy.kotlin.codegen.core.withBlock
 import software.amazon.smithy.kotlin.codegen.core.withInlineBlock
 import software.amazon.smithy.kotlin.codegen.lang.KotlinTypes
 import software.amazon.smithy.kotlin.codegen.model.getTrait
+import software.amazon.smithy.kotlin.codegen.service.contraints.ConstraintGenerator
+import software.amazon.smithy.kotlin.codegen.service.contraints.ConstraintUtilsGenerator
 import software.amazon.smithy.model.shapes.OperationShape
 import software.amazon.smithy.model.traits.AuthTrait
 import software.amazon.smithy.model.traits.HttpBearerAuthTrait
@@ -93,7 +95,7 @@ internal class KtorStubGenerator(
     }
 
     private fun renderLogging() {
-        delegator.useFileWriter("Logging.kt", "${ctx.settings.pkg.name}.utils") { writer ->
+        delegator.useFileWriter("Logging.kt", "$pkgName.utils") { writer ->
 
             writer.withBlock("internal fun #T.configureLogging() {", "}", RuntimeTypes.KtorServerCore.Application) {
                 withBlock(
@@ -167,20 +169,20 @@ internal class KtorStubGenerator(
 
     // Generates `Authentication.kt` with Authenticator interface + configureSecurity().
     override fun renderAuthModule() {
-        delegator.useFileWriter("UserPrincipal.kt", "${ctx.settings.pkg.name}.auth") { writer ->
+        delegator.useFileWriter("UserPrincipal.kt", "$pkgName.auth") { writer ->
             writer.withBlock("public data class UserPrincipal(", ")") {
                 write("val user: String")
             }
         }
 
-        delegator.useFileWriter("Validation.kt", "${ctx.settings.pkg.name}.auth") { writer ->
+        delegator.useFileWriter("Validation.kt", "$pkgName.auth") { writer ->
             writer.withBlock("public fun bearerValidation(token: String): UserPrincipal? {", "}") {
                 write("// TODO: implement me")
                 write("if (true) return UserPrincipal(#S) else return null", "Authenticated User")
             }
         }
 
-        delegator.useFileWriter("Authentication.kt", "${ctx.settings.pkg.name}.auth") { writer ->
+        delegator.useFileWriter("Authentication.kt", "$pkgName.auth") { writer ->
             writer.withBlock("internal fun #T.configureAuthentication() {", "}", RuntimeTypes.KtorServerCore.Application) {
                 write("")
                 withBlock(
@@ -201,20 +203,23 @@ internal class KtorStubGenerator(
 
     // For every operation request structure, create a `validate()` function file.
     override fun renderConstraintValidators() {
+        ConstraintUtilsGenerator(ctx, delegator).render()
+        operations.forEach { operation -> ConstraintGenerator(ctx, operation, delegator).render() }
     }
 
     // Writes `Routing.kt` that maps Smithy operations → Ktor routes.
     override fun renderRouting() {
         val contentType = ContentType.fromServiceShape(serviceShape)
 
-        delegator.useFileWriter("Routing.kt", ctx.settings.pkg.name) { writer ->
+        delegator.useFileWriter("Routing.kt", pkgName) { writer ->
 
             operations.forEach { shape ->
-                writer.addImport("${ctx.settings.pkg.name}.serde", "${shape.id.name}OperationDeserializer")
-                writer.addImport("${ctx.settings.pkg.name}.serde", "${shape.id.name}OperationSerializer")
-                writer.addImport("${ctx.settings.pkg.name}.model", "${shape.id.name}Request")
-                writer.addImport("${ctx.settings.pkg.name}.model", "${shape.id.name}Response")
-                writer.addImport("${ctx.settings.pkg.name}.operations", "handle${shape.id.name}Request")
+                writer.addImport("$pkgName.serde", "${shape.id.name}OperationDeserializer")
+                writer.addImport("$pkgName.serde", "${shape.id.name}OperationSerializer")
+                writer.addImport("$pkgName.constraints", "check${shape.id.name}RequestConstraint")
+                writer.addImport("$pkgName.model", "${shape.id.name}Request")
+                writer.addImport("$pkgName.model", "${shape.id.name}Response")
+                writer.addImport("$pkgName.operations", "handle${shape.id.name}Request")
             }
 
             writer.withBlock("internal fun #T.configureRouting(): Unit {", "}", RuntimeTypes.KtorServerCore.Application) {
@@ -271,6 +276,11 @@ internal class KtorStubGenerator(
                                                     "Malformed CBOR input",
                                                 )
                                             }
+                                            write(
+                                                "try { check${shape.id.name}RequestConstraint(requestObj) } catch (ex: Exception) { throw #T(ex?.message ?: #S, ex) }",
+                                                RuntimeTypes.KtorServerCore.BadRequestException,
+                                                "Error while validating constraints",
+                                            )
                                             write("val responseObj = handle${shape.id.name}Request(requestObj)")
                                             write("val serializer = ${shape.id.name}OperationSerializer()")
                                             withBlock(
@@ -356,7 +366,7 @@ internal class KtorStubGenerator(
     }
 
     private fun renderErrorHandler() {
-        delegator.useFileWriter("ErrorHandler.kt", "${ctx.settings.pkg.name}.plugins") { writer ->
+        delegator.useFileWriter("ErrorHandler.kt", "$pkgName.plugins") { writer ->
             writer.write("@#T", RuntimeTypes.KotlinxCborSerde.Serializable)
                 .write("private data class ErrorPayload(val code: Int, val message: String)")
                 .write("")
@@ -386,7 +396,7 @@ internal class KtorStubGenerator(
                     write("val acceptsCbor = request.#T().any { it.value == #S }", RuntimeTypes.KtorServerRouting.requestAcceptItems, "application/cbor")
                     write("val acceptsJson = request.#T().any { it.value == #S }", RuntimeTypes.KtorServerRouting.requestAcceptItems, "application/json")
                     write("")
-                    write("val log = #T.getLogger(#S)", RuntimeTypes.KtorLoggingSlf4j.LoggerFactory, ctx.settings.pkg.name)
+                    write("val log = #T.getLogger(#S)", RuntimeTypes.KtorLoggingSlf4j.LoggerFactory, pkgName)
                     write("log.info(#S)", "Route Error Message: \${envelope.msg}")
                     write("")
                     withBlock("when {", "}") {
@@ -427,6 +437,16 @@ internal class KtorStubGenerator(
                             write("call.respondEnvelope( ErrorEnvelope(status.value, message), status )")
                         }
                         write("")
+                        withBlock("status(#T.NotFound) { call, status ->", "}", RuntimeTypes.KtorServerHttp.HttpStatusCode) {
+                            write("val message = #S", "Resource not found")
+                            write("call.respondEnvelope( ErrorEnvelope(status.value, message), status )")
+                        }
+                        write("")
+                        withBlock("status(#T.MethodNotAllowed) { call, status ->", "}", RuntimeTypes.KtorServerHttp.HttpStatusCode) {
+                            write("val message = #S", "Method not allowed for this resource")
+                            write("call.respondEnvelope( ErrorEnvelope(status.value, message), status )")
+                        }
+                        write("")
                         withBlock("#T<Throwable> { call, cause ->", "}", RuntimeTypes.KtorServerStatusPage.exception) {
                             withBlock("val status = when (cause) {", "}") {
                                 write(
@@ -456,7 +476,7 @@ internal class KtorStubGenerator(
     }
 
     private fun renderContentTypeGuard() {
-        delegator.useFileWriter("ContentTypeGuard.kt", "${ctx.settings.pkg.name}.plugins") { writer ->
+        delegator.useFileWriter("ContentTypeGuard.kt", "$pkgName.plugins") { writer ->
 
             writer.withBlock("private fun #T.hasBody(): Boolean {", "}", RuntimeTypes.KtorServerRouting.requestApplicationRequest) {
                 write(
@@ -569,9 +589,9 @@ internal class KtorStubGenerator(
     override fun renderPerOperationHandlers() {
         operations.forEach { shape ->
             val name = shape.id.name
-            delegator.useFileWriter("${name}Operation.kt", "${ctx.settings.pkg.name}.operations") { writer ->
-                writer.addImport("${ctx.settings.pkg.name}.model", "${shape.id.name}Request")
-                writer.addImport("${ctx.settings.pkg.name}.model", "${shape.id.name}Response")
+            delegator.useFileWriter("${name}Operation.kt", "$pkgName.operations") { writer ->
+                writer.addImport("$pkgName.model", "${shape.id.name}Request")
+                writer.addImport("$pkgName.model", "${shape.id.name}Response")
 
                 writer.withBlock("public fun handle${name}Request(req: ${name}Request): ${name}Response {", "}") {
                     write("// TODO: implement me")

codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/service/ServiceTypes.kt

@@ -67,4 +67,14 @@ class ServiceTypes(val pkgName: String) {
         name = "AcceptTypeGuard"
         namespace = "$pkgName.plugins"
     }
+
+    val sizeOf = buildSymbol {
+        name = "sizeOf"
+        namespace = "$pkgName.constraints"
+    }
+
+    val hasAllUniqueElements = buildSymbol {
+        name = "hasAllUniqueElements"
+        namespace = "$pkgName.constraints"
+    }
 }

codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/service/contraints/AbstractConstraintTraitGenerator.kt

@@ -0,0 +1,5 @@
+package software.amazon.smithy.kotlin.codegen.service.contraints
+
+internal abstract class AbstractConstraintTraitGenerator {
+    abstract fun render()
+}

codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/service/contraints/ConstraintGenerator.kt

@@ -0,0 +1,62 @@
+package software.amazon.smithy.kotlin.codegen.service.contraints
+
+import software.amazon.smithy.kotlin.codegen.core.GenerationContext
+import software.amazon.smithy.kotlin.codegen.core.KotlinDelegator
+import software.amazon.smithy.kotlin.codegen.core.KotlinWriter
+import software.amazon.smithy.kotlin.codegen.core.withBlock
+import software.amazon.smithy.model.shapes.MemberShape
+import software.amazon.smithy.model.shapes.OperationShape
+import software.amazon.smithy.model.shapes.StructureShape
+import software.amazon.smithy.model.traits.RequiredTrait
+import kotlin.collections.iterator
+
+internal class ConstraintGenerator(
+    val ctx: GenerationContext,
+    val operation: OperationShape,
+    val delegator: KotlinDelegator,
+) {
+    val inputShape = ctx.model.expectShape(operation.input.get()) as StructureShape
+    val inputMembers = inputShape.allMembers
+
+    val opName = operation.id.name
+    val pkgName = ctx.settings.pkg.name
+
+    fun render() {
+        renderRequestConstraintsValidation()
+    }
+    private fun generateConstraintValidations(prefix: String, memberShape: MemberShape, writer: KotlinWriter) {
+        val targetShape = ctx.model.expectShape(memberShape.target)
+
+        val memberName = memberShape.memberName
+        val memberAndTargetTraits = memberShape.allTraits + targetShape.allTraits
+
+        for (memberTrait in memberAndTargetTraits.values) {
+            val traitGenerator = getTraitGeneratorFromTrait(prefix, memberName, memberTrait, pkgName, writer)
+            if (memberTrait !is RequiredTrait) {
+                writer.write("if ($prefix$memberName == null) { return }")
+            }
+            traitGenerator?.render()
+        }
+
+        for (member in targetShape.allMembers) {
+            val newMemberPrefix = "${targetShape.id.name}".replaceFirstChar { it.lowercase() }
+            writer.withBlock("if ($prefix$memberName != null) {", "}") {
+                withBlock("for ($newMemberPrefix${member.key} in $prefix$memberName) {", "}") {
+                    call { generateConstraintValidations(newMemberPrefix, member.value, writer) }
+                }
+            }
+        }
+    }
+
+    private fun renderRequestConstraintsValidation() {
+        delegator.useFileWriter("${opName}RequestConstraints.kt", "$pkgName.constraints") { writer ->
+            writer.addImport("$pkgName.model", "${operation.id.name}Request")
+
+            writer.withBlock("public fun check${opName}RequestConstraint(data: ${opName}Request) {", "}") {
+                for (memberShape in inputMembers.values) {
+                    generateConstraintValidations("data.", memberShape, writer)
+                }
+            }
+        }
+    }
+}

codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/service/contraints/ConstraintUtilsGenerator.kt

@@ -0,0 +1,75 @@
+package software.amazon.smithy.kotlin.codegen.service.contraints
+
+import software.amazon.smithy.kotlin.codegen.core.GenerationContext
+import software.amazon.smithy.kotlin.codegen.core.KotlinDelegator
+import software.amazon.smithy.kotlin.codegen.core.KotlinWriter
+import software.amazon.smithy.kotlin.codegen.core.withBlock
+
+internal class ConstraintUtilsGenerator(
+    val ctx: GenerationContext,
+    val delegator: KotlinDelegator,
+) {
+    val pkgName = ctx.settings.pkg.name
+
+    fun render() {
+        delegator.useFileWriter("utils.kt", "$pkgName.constraints") { writer ->
+            renderLengthTraitUtils(writer)
+
+            writer.write("")
+            renderUniqueItemsTraitUtils(writer)
+        }
+    }
+
+    private fun renderLengthTraitUtils(writer: KotlinWriter) {
+        writer.withBlock("internal fun sizeOf(value: Any?): Long = when (value) {", "}") {
+            write("is Collection<*> -> value.size.toLong()")
+            write("is Array<*> -> value.size.toLong()")
+            write("is Map<*, *> -> value.size.toLong()")
+            write("is String -> value.codePointCount(0, value.length).toLong()")
+            write("is ByteArray -> value.size.toLong()")
+            withBlock("else -> {", "}") {
+                write("val typeName = value?.javaClass?.simpleName ?: #S", "null")
+                write("throw IllegalArgumentException( #S )", "sizeOf does not support \${typeName} type")
+            }
+        }
+    }
+
+    private fun renderUniqueItemsTraitUtils(writer: KotlinWriter) {
+        writer.withBlock("internal fun hasAllUniqueElements(elements: List<Any?>): Boolean {", "}") {
+            withBlock("class Wrapped(private val v: Any?) {", "}") {
+                withBlock("override fun equals(other: Any?): Boolean {", "}") {
+                    write("if (other !is Wrapped) return false")
+                    write("if (v?.javaClass != other.v?.javaClass) return false")
+                    withBlock("return when (v) {", "}") {
+                        write("null -> true")
+                        write("is String,")
+                        write("is Boolean,")
+                        write("is java.time.Instant,")
+                        write("is Number -> v == other.v")
+                        write("is ByteArray -> v.contentEquals(other.v as ByteArray)")
+                        withBlock("is List<*> -> {", "}") {
+                            write("val o = other.v as List<*>")
+                            write("v.size == o.size &&  v.indices.all { i -> Wrapped(v[i]) == Wrapped(o[i]) }")
+                        }
+                        withBlock("is Map<*, *>  -> {", "}") {
+                            write("val o = other.v as Map<*, *>")
+                            write("v.size == o.size && v.all { (k, value) -> o.containsKey(k) && Wrapped(value) == Wrapped(o[k]) }")
+                        }
+                        write("else -> v == other.v")
+                    }
+                }
+                withBlock("override fun hashCode(): Int = when (v) {", "}") {
+                    write("null -> 0")
+                    write("is ByteArray -> v.contentHashCode()")
+                    write("is List<*> -> v.fold(1) { acc, e -> 31 * acc + Wrapped(e).hashCode() }")
+                    write("is Map<*, *> -> v.entries.fold(1) { acc, (k, e) -> 31 * acc + Wrapped(k).hashCode() xor Wrapped(e).hashCode() }")
+                    write("else -> v.hashCode()")
+                }
+            }
+            write("")
+            write("val seen = HashSet<Wrapped>(elements.size)")
+            write("for (e in elements) if (!seen.add(Wrapped(e))) return false")
+            write("return true")
+        }
+    }
+}

codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/service/contraints/LengthConstraintGenerator.kt

@@ -0,0 +1,21 @@
+package software.amazon.smithy.kotlin.codegen.service.contraints
+
+import software.amazon.smithy.kotlin.codegen.core.KotlinWriter
+import software.amazon.smithy.kotlin.codegen.service.ServiceTypes
+import software.amazon.smithy.model.traits.LengthTrait
+
+internal class LengthConstraintGenerator(val memberPrefix: String, val memberName: String, val trait: LengthTrait, val pkgName: String, val writer: KotlinWriter) : AbstractConstraintTraitGenerator() {
+    override fun render() {
+        val min = trait.min.orElse(null)
+        val max = trait.max.orElse(null)
+        val member = "$memberPrefix$memberName"
+
+        if (max != null && min != null) {
+            writer.write("require(#T($member) in $min..$max) { #S }", ServiceTypes(pkgName).sizeOf, "The size of `$memberName` must be between $min and $max (inclusive)")
+        } else if (max != null) {
+            writer.write("require(#T($member) <= $max) { #S }", ServiceTypes(pkgName).sizeOf, "The size of `$memberName` must be less than or equal to $max")
+        } else {
+            writer.write("require(#T($member) >= $min) { #S }", ServiceTypes(pkgName).sizeOf, "The size of `$memberName` must be greater than or equal to $min")
+        }
+    }
+}

codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/service/contraints/PatternConstraintGenerator.kt

@@ -0,0 +1,12 @@
+package software.amazon.smithy.kotlin.codegen.service.contraints
+
+import software.amazon.smithy.kotlin.codegen.core.KotlinWriter
+import software.amazon.smithy.model.traits.PatternTrait
+
+internal class PatternConstraintGenerator(val memberPrefix: String, val memberName: String, val trait: PatternTrait, val pkgName: String, val writer: KotlinWriter) : AbstractConstraintTraitGenerator() {
+    override fun render() {
+        val member = "$memberPrefix$memberName"
+
+        writer.write("require(Regex(#S).containsMatchIn($member)) { #S }", trait.pattern.toString(), "Value `\${$member}` does not match required pattern: `${trait.pattern}`")
+    }
+}

codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/service/contraints/RangeConstraintGenerator.kt

@@ -0,0 +1,20 @@
+package software.amazon.smithy.kotlin.codegen.service.contraints
+
+import software.amazon.smithy.kotlin.codegen.core.KotlinWriter
+import software.amazon.smithy.model.traits.RangeTrait
+
+internal class RangeConstraintGenerator(val memberPrefix: String, val memberName: String, val trait: RangeTrait, val pkgName: String, val writer: KotlinWriter) : AbstractConstraintTraitGenerator() {
+    override fun render() {
+        val min = trait.min.orElse(null)
+        val max = trait.max.orElse(null)
+        val member = "$memberPrefix$memberName"
+
+        if (max != null && min != null) {
+            writer.write("require($member in $min..$max) { #S }", "`$memberName` must be between $min and $max (inclusive)")
+        } else if (max != null) {
+            writer.write("require($member <= $max) { #S }", "`$memberName` must be less than or equal to $max")
+        } else {
+            writer.write("require($member >= $min) { #S }", "`$memberName` must be greater than or equal to $min")
+        }
+    }
+}

codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/service/contraints/RequiredConstraintGenerator.kt

@@ -0,0 +1,11 @@
+package software.amazon.smithy.kotlin.codegen.service.contraints
+
+import software.amazon.smithy.kotlin.codegen.core.KotlinWriter
+import software.amazon.smithy.model.traits.RequiredTrait
+
+internal class RequiredConstraintGenerator(val memberPrefix: String, val memberName: String, val trait: RequiredTrait, val pkgName: String, val writer: KotlinWriter) : AbstractConstraintTraitGenerator() {
+    override fun render() {
+        val member = "$memberPrefix$memberName"
+        writer.write("require($member != null) { #S }", "`$memberName` must be provided")
+    }
+}