From accebab83f6ad36723a899631aa36e81f937d56e Mon Sep 17 00:00:00 2001
From: Lukas Ruegner <ruegnerlukas@gmail.com>
Date: Sun, 17 Nov 2024 16:14:41 +0100
Subject: [PATCH] store security-schemas in components as configured

---
 .../ktorswaggerui/examples/Authentication.kt  | 11 +++++++++
 .../builder/schema/SchemaContextImpl.kt       | 23 ++++---------------
 .../builder/OperationBuilderTest.kt           |  2 +-
 3 files changed, 17 insertions(+), 19 deletions(-)

diff --git a/ktor-swagger-ui-examples/src/main/kotlin/io/github/smiley4/ktorswaggerui/examples/Authentication.kt b/ktor-swagger-ui-examples/src/main/kotlin/io/github/smiley4/ktorswaggerui/examples/Authentication.kt
index 5ae0e50..1faa78e 100644
--- a/ktor-swagger-ui-examples/src/main/kotlin/io/github/smiley4/ktorswaggerui/examples/Authentication.kt
+++ b/ktor-swagger-ui-examples/src/main/kotlin/io/github/smiley4/ktorswaggerui/examples/Authentication.kt
@@ -42,6 +42,7 @@ private fun Application.myModule() {
 
     // Install and configure the "SwaggerUI"-Plugin
     install(SwaggerUI) {
+        schemas {  }
         security {
             // configure a basic-auth security scheme
             securityScheme("MySecurityScheme") {
@@ -53,6 +54,7 @@ private fun Application.myModule() {
             // if no other response is documented for "401 Unauthorized", this information is used instead
             defaultUnauthorizedResponse {
                 description = "Username or password is invalid"
+                body<AuthRequired>()
             }
         }
     }
@@ -74,6 +76,12 @@ private fun Application.myModule() {
             }) {
                 call.respondText("Hello World!")
             }
+
+            get("protected2", {
+                // response for "401 Unauthorized" is automatically added if configured in the plugin-config and not specified otherwise
+            }) {
+                call.respondText("Hello World!")
+            }
         }
 
         // route is not in an "authenticate"-block but "protected"-flag is set (e.g. because is it protected by an external reverse-proxy
@@ -96,3 +104,6 @@ private fun Application.myModule() {
     }
 
 }
+
+
+class AuthRequired(val message: String)
\ No newline at end of file
diff --git a/ktor-swagger-ui/src/main/kotlin/io/github/smiley4/ktorswaggerui/builder/schema/SchemaContextImpl.kt b/ktor-swagger-ui/src/main/kotlin/io/github/smiley4/ktorswaggerui/builder/schema/SchemaContextImpl.kt
index ddce850..e1cfac5 100644
--- a/ktor-swagger-ui/src/main/kotlin/io/github/smiley4/ktorswaggerui/builder/schema/SchemaContextImpl.kt
+++ b/ktor-swagger-ui/src/main/kotlin/io/github/smiley4/ktorswaggerui/builder/schema/SchemaContextImpl.kt
@@ -15,11 +15,14 @@ class SchemaContextImpl(private val schemaConfig: SchemaConfigData) : SchemaCont
 
     fun addGlobal(config: SchemaConfigData) {
         config.securitySchemas.forEach { typeDescriptor ->
-            val schema = collapseRootRef(generateSchema(typeDescriptor))
+            val schema = generateSchema(typeDescriptor)
             rootSchemas[typeDescriptor] = schema.swagger
+            schema.componentSchemas.forEach { (k, v) ->
+                componentSchemas[k] = v
+            }
         }
         config.schemas.forEach { (schemaId, typeDescriptor) ->
-            val schema = collapseRootRef(generateSchema(typeDescriptor))
+            val schema = generateSchema(typeDescriptor)
             componentSchemas[schemaId] = schema.swagger
             schema.componentSchemas.forEach { (k, v) ->
                 componentSchemas[k] = v
@@ -27,22 +30,6 @@ class SchemaContextImpl(private val schemaConfig: SchemaConfigData) : SchemaCont
         }
     }
 
-    private fun collapseRootRef(schema: CompiledSwaggerSchema): CompiledSwaggerSchema {
-        if (schema.swagger.`$ref` == null) {
-            return schema
-        } else {
-            val referencedSchemaId = schema.swagger.`$ref`!!.replace("#/components/schemas/", "")
-            val referencedSchema = schema.componentSchemas[referencedSchemaId]!!
-            return CompiledSwaggerSchema(
-                typeData = schema.typeData,
-                swagger = referencedSchema,
-                componentSchemas = schema.componentSchemas.toMutableMap().also {
-                    it.remove(referencedSchemaId)
-                }
-            )
-        }
-    }
-
     fun add(routes: Collection<RouteMeta>) {
         collectTypeDescriptor(routes).forEach { typeDescriptor ->
             val schema = generateSchema(typeDescriptor)
diff --git a/ktor-swagger-ui/src/test/kotlin/io/github/smiley4/ktorswaggerui/builder/OperationBuilderTest.kt b/ktor-swagger-ui/src/test/kotlin/io/github/smiley4/ktorswaggerui/builder/OperationBuilderTest.kt
index e39e37c..ef51e82 100644
--- a/ktor-swagger-ui/src/test/kotlin/io/github/smiley4/ktorswaggerui/builder/OperationBuilderTest.kt
+++ b/ktor-swagger-ui/src/test/kotlin/io/github/smiley4/ktorswaggerui/builder/OperationBuilderTest.kt
@@ -754,7 +754,7 @@ class OperationBuilderTest : StringSpec({
                         ?.also { response ->
                             val content = response.content["application/json"]
                             content.shouldNotBeNull()
-                            content.schema.title shouldBe "SimpleObject"
+                            content.schema.`$ref` shouldBe "#/components/schemas/${SimpleObject::class.qualifiedName}"
                             val example = content.examples["Example 1"]
                             example.shouldNotBeNull()
                             example.value shouldBeEqual SimpleObject(text = "Some text", number = 123)