Skip to content

feat(core): security scheme extraction and auth-aware code generation #31

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
halotukozak wants to merge 1 commit into
base: refactor/issue-model
Choose a base branch
from
Open

feat(core): security scheme extraction and auth-aware code generation #31

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
156 changes: 129 additions & 27 deletions core/src/main/kotlin/com/avsystem/justworks/core/gen/ApiClientBaseGenerator.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
package com.avsystem.justworks.core.gen

import com.avsystem.justworks.core.model.ApiKeyLocation
import com.avsystem.justworks.core.model.SecurityScheme
import com.squareup.kotlinpoet.CodeBlock
import com.squareup.kotlinpoet.ContextParameter
import com.squareup.kotlinpoet.ExperimentalKotlinPoetApi
Expand Down Expand Up @@ -30,7 +32,7 @@ object ApiClientBaseGenerator {
private const val SUCCESS_BODY = "successBody"
private const val SERIALIZERS_MODULE_PARAM = "serializersModule"

fun generate(): FileSpec {
fun generate(securitySchemes: List<SecurityScheme>): FileSpec {
val t = TypeVariableName("T").copy(reified = true)

return FileSpec
Expand All @@ -39,7 +41,7 @@ object ApiClientBaseGenerator {
.addFunction(buildMapToResult(t))
.addFunction(buildToResult(t))
.addFunction(buildToEmptyResult())
.addType(buildApiClientBaseClass())
.addType(buildApiClientBaseClass(securitySchemes))
.build()
}

Expand Down Expand Up @@ -103,26 +105,33 @@ object ApiClientBaseGenerator {
.addStatement("return %L { Unit }", MAP_TO_RESULT)
.build()

private fun buildApiClientBaseClass(): TypeSpec {
private fun buildApiClientBaseClass(securitySchemes: List<SecurityScheme>): TypeSpec {
val tokenType = LambdaTypeName.get(returnType = STRING)
val authParams = buildAuthConstructorParams(securitySchemes)

val constructor = FunSpec
val constructorBuilder = FunSpec
.constructorBuilder()
.addParameter(BASE_URL, STRING)
.addParameter(TOKEN, tokenType)
.build()

val propertySpecs = mutableListOf<PropertySpec>()

val baseUrlProp = PropertySpec
.builder(BASE_URL, STRING)
.initializer(BASE_URL)
.addModifiers(KModifier.PROTECTED)
.build()
propertySpecs.add(baseUrlProp)

val tokenProp = PropertySpec
.builder(TOKEN, tokenType)
.initializer(TOKEN)
.addModifiers(KModifier.PRIVATE)
.build()
for ((paramName, _) in authParams) {
constructorBuilder.addParameter(paramName, tokenType)
propertySpecs.add(
PropertySpec
.builder(paramName, tokenType)
.initializer(paramName)
.addModifiers(KModifier.PRIVATE)
.build(),
)
}

val clientProp = PropertySpec
.builder(CLIENT, HTTP_CLIENT)
Expand All @@ -135,32 +144,125 @@ object ApiClientBaseGenerator {
.addStatement("$CLIENT.close()")
.build()

return TypeSpec
val classBuilder = TypeSpec
.classBuilder(API_CLIENT_BASE)
.addModifiers(KModifier.ABSTRACT)
.addSuperinterface(CLOSEABLE)
.primaryConstructor(constructor)
.addProperty(baseUrlProp)
.addProperty(tokenProp)
.primaryConstructor(constructorBuilder.build())

for (prop in propertySpecs) {
classBuilder.addProperty(prop)
}

return classBuilder
.addProperty(clientProp)
.addFunction(closeFun)
.addFunction(buildApplyAuth())
.addFunction(buildApplyAuth(securitySchemes))
.addFunction(buildSafeCall())
.addFunction(buildCreateHttpClient())
.build()
}

private fun buildApplyAuth(): FunSpec = FunSpec
.builder(APPLY_AUTH)
.addModifiers(KModifier.PROTECTED)
.receiver(HTTP_REQUEST_BUILDER)
.beginControlFlow("%M", HEADERS_FUN)
.addStatement(
"append(%T.Authorization, %P)",
HTTP_HEADERS,
CodeBlock.of($$"Bearer ${'$'}{$$TOKEN()}"),
).endControlFlow()
.build()
/**
* Builds the list of auth-related constructor parameter names based on security schemes.
* Returns pairs of (paramName, schemeType) for each scheme.
*/
internal fun buildAuthConstructorParams(securitySchemes: List<SecurityScheme>): List<Pair<String, SecurityScheme>> =
securitySchemes.flatMap { scheme ->
when (scheme) {
is SecurityScheme.Bearer -> {
val isSingleBearer =
securitySchemes.size == 1 && securitySchemes.first() is SecurityScheme.Bearer

val paramName = if (isSingleBearer) TOKEN else "${scheme.name.toCamelCase()}Token"
listOf(paramName to scheme)
}

is SecurityScheme.ApiKey -> {
listOf("${scheme.name.toCamelCase()}Key" to scheme)
}

is SecurityScheme.Basic -> {
listOf(
"${scheme.name.toCamelCase()}Username" to scheme,
"${scheme.name.toCamelCase()}Password" to scheme,
)
}
}
}

private fun buildApplyAuth(securitySchemes: List<SecurityScheme>): FunSpec {
val builder = FunSpec
.builder(APPLY_AUTH)
.addModifiers(KModifier.PROTECTED)
.receiver(HTTP_REQUEST_BUILDER)

if (securitySchemes.isEmpty()) return builder.build()

val headerSchemes = securitySchemes.filter {
it is SecurityScheme.Bearer ||
it is SecurityScheme.Basic ||
(it is SecurityScheme.ApiKey && it.location == ApiKeyLocation.HEADER)
}
val querySchemes = securitySchemes
.filterIsInstance<SecurityScheme.ApiKey>()
.filter { it.location == ApiKeyLocation.QUERY }

if (headerSchemes.isNotEmpty()) {
builder.beginControlFlow("%M", HEADERS_FUN)
for (scheme in headerSchemes) {
when (scheme) {
is SecurityScheme.Bearer -> {
val isSingleBearer =
securitySchemes.size == 1 && securitySchemes.first() is SecurityScheme.Bearer

val paramName = if (isSingleBearer) TOKEN else "${scheme.name.toCamelCase()}Token"
builder.addStatement(
"append(%T.Authorization, %P)",
HTTP_HEADERS,
CodeBlock.of("Bearer \${$paramName()}"),
)
}

is SecurityScheme.Basic -> {
val usernameParam = "${scheme.name.toCamelCase()}Username"
val passwordParam = "${scheme.name.toCamelCase()}Password"
builder.addStatement(
"append(%T.Authorization, %P)",
HTTP_HEADERS,
CodeBlock.of(
"Basic \${%T.getEncoder().encodeToString(\"${'$'}{$usernameParam()}:${'$'}{$passwordParam()}\".toByteArray())}",
BASE64_CLASS,
),
)
Comment on lines +211 to +237
Copy link

Copilot AI Mar 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When multiple Bearer/Basic schemes are present, this implementation will append(HttpHeaders.Authorization, ...) multiple times, producing multiple Authorization headers on a single request. Many servers reject or mis-handle that header; consider enforcing at most one Authorization-based scheme, or generating a single chosen header based on security requirements rather than applying all schemes unconditionally.

Copilot uses AI. Check for mistakes.
}

is SecurityScheme.ApiKey -> {
val paramName = "${scheme.name.toCamelCase()}Key"
builder.addStatement(
"append(%S, $paramName())",
scheme.parameterName,
)
}
}
}
builder.endControlFlow()
}

if (querySchemes.isNotEmpty()) {
builder.beginControlFlow("url")
for (scheme in querySchemes) {
val paramName = "${scheme.name.toCamelCase()}Key"
builder.addStatement(
"parameters.append(%S, $paramName())",
scheme.parameterName,
)
}
builder.endControlFlow()
}

return builder.build()
}

private fun buildSafeCall(): FunSpec = FunSpec
.builder(SAFE_CALL)
Expand Down
29 changes: 19 additions & 10 deletions core/src/main/kotlin/com/avsystem/justworks/core/gen/ClientGenerator.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import com.avsystem.justworks.core.model.Endpoint
import com.avsystem.justworks.core.model.HttpMethod
import com.avsystem.justworks.core.model.Parameter
import com.avsystem.justworks.core.model.ParameterLocation
import com.avsystem.justworks.core.model.SecurityScheme
import com.avsystem.justworks.core.model.TypeRef
import com.squareup.kotlinpoet.ClassName
import com.squareup.kotlinpoet.CodeBlock
Expand Down Expand Up @@ -34,13 +35,16 @@ private const val API_SUFFIX = "Api"
class ClientGenerator(private val apiPackage: String, private val modelPackage: String) {
fun generate(spec: ApiSpec, hasPolymorphicTypes: Boolean = false): List<FileSpec> {
val grouped = spec.endpoints.groupBy { it.tags.firstOrNull() ?: DEFAULT_TAG }
return grouped.map { (tag, endpoints) -> generateClientFile(tag, endpoints, hasPolymorphicTypes) }
return grouped.map { (tag, endpoints) ->
generateClientFile(tag, endpoints, hasPolymorphicTypes, spec.securitySchemes)
}
}

private fun generateClientFile(
tag: String,
endpoints: List<Endpoint>,
hasPolymorphicTypes: Boolean = false,
securitySchemes: List<SecurityScheme>,
): FileSpec {
val className = ClassName(apiPackage, "${tag.toPascalCase()}$API_SUFFIX")

Expand All @@ -52,25 +56,30 @@ class ClientGenerator(private val apiPackage: String, private val modelPackage:
}

val tokenType = LambdaTypeName.get(returnType = STRING)
val authParams = ApiClientBaseGenerator.buildAuthConstructorParams(securitySchemes)

val primaryConstructor = FunSpec
val constructorBuilder = FunSpec
.constructorBuilder()
.addParameter(BASE_URL, STRING)
.addParameter(TOKEN, tokenType)
.build()

val classBuilder = TypeSpec
.classBuilder(className)
.superclass(API_CLIENT_BASE)
.addSuperclassConstructorParameter(BASE_URL)

for ((paramName, _) in authParams) {
constructorBuilder.addParameter(paramName, tokenType)
classBuilder.addSuperclassConstructorParameter(paramName)
}

val httpClientProperty = PropertySpec
.builder(CLIENT, HTTP_CLIENT)
.addModifiers(KModifier.OVERRIDE, KModifier.PROTECTED)
.initializer(clientInitializer)
.build()

val classBuilder = TypeSpec
.classBuilder(className)
.superclass(API_CLIENT_BASE)
.addSuperclassConstructorParameter(BASE_URL)
.addSuperclassConstructorParameter(TOKEN)
.primaryConstructor(primaryConstructor)
classBuilder
.primaryConstructor(constructorBuilder.build())
.addProperty(httpClientProperty)

classBuilder.addFunctions(endpoints.map(::generateEndpointFunction))
Expand Down
8 changes: 5 additions & 3 deletions core/src/main/kotlin/com/avsystem/justworks/core/gen/CodeGenerator.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ object CodeGenerator {
spec: ApiSpec,
modelPackage: String,
apiPackage: String,
outputDir: File
outputDir: File,
): Result {
val modelFiles = ModelGenerator(modelPackage).generate(spec)
modelFiles.forEach { it.writeTo(outputDir) }
Expand All @@ -27,8 +27,10 @@ object CodeGenerator {
return Result(modelFiles.size, clientFiles.size)
}

fun generateSharedTypes(outputDir: File): Int {
val files = ApiResponseGenerator.generate() + ApiClientBaseGenerator.generate()
fun generateSharedTypes(outputDir: File, specs: List<ApiSpec> = emptyList()): Int {
val securitySchemes = specs.flatMap { it.securitySchemes }

val files = ApiResponseGenerator.generate() + ApiClientBaseGenerator.generate(securitySchemes)
Comment on lines +31 to +33
Copy link

Copilot AI Mar 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

generateSharedTypes builds ApiClientBase from the combined securitySchemes of all specs, but ClientGenerator builds each client constructor from its own spec’s schemes. In a multi-spec project where scheme sets differ, generated clients won’t match the shared ApiClientBase constructor and compilation will fail. Consider keeping ApiClientBase constructor stable (e.g., defaults/optional providers), generating ApiClientBase per spec instead of as a shared type, or ensuring both shared types and clients are generated from the same unified scheme set.

Suggested change
val securitySchemes = specs.flatMap { it.securitySchemes }
val files = ApiResponseGenerator.generate() + ApiClientBaseGenerator.generate(securitySchemes)
val files = ApiResponseGenerator.generate() + ApiClientBaseGenerator.generate()

Copilot uses AI. Check for mistakes.
files.forEach { it.writeTo(outputDir) }
return files.size
}
Expand Down
1 change: 1 addition & 0 deletions core/src/main/kotlin/com/avsystem/justworks/core/gen/Names.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,7 @@ val HTTP_SUCCESS = ClassName("com.avsystem.justworks", "HttpSuccess")
// Kotlin stdlib
// ============================================================================

val BASE64_CLASS = ClassName("java.util", "Base64")
val CLOSEABLE = ClassName("java.io", "Closeable")
val IO_EXCEPTION = ClassName("java.io", "IOException")
val HTTP_REQUEST_TIMEOUT_EXCEPTION = ClassName("io.ktor.client.plugins", "HttpRequestTimeoutException")
Expand Down
17 changes: 17 additions & 0 deletions core/src/main/kotlin/com/avsystem/justworks/core/model/ApiSpec.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,12 +7,29 @@ package com.avsystem.justworks.core.model
* code generators. Bridges the raw Swagger Parser OAS model and the generated
* Kotlin client/model source files.
*/
sealed interface SecurityScheme {
val name: String

data class Bearer(override val name: String) : SecurityScheme

data class ApiKey(
override val name: String,
val parameterName: String,
val location: ApiKeyLocation,
) : SecurityScheme

data class Basic(override val name: String) : SecurityScheme
}

enum class ApiKeyLocation { HEADER, QUERY }

data class ApiSpec(
val title: String,
val version: String,
val endpoints: List<Endpoint>,
val schemas: List<SchemaModel>,
val enums: List<EnumModel>,
val securitySchemes: List<SecurityScheme>,
)

data class Endpoint(
Expand Down
Loading
Loading