feat: community auth (#13175)

Co-authored-by: Parker Mossman <[email protected]>

airbyte-commons-auth/src/main/kotlin/io/airbyte/commons/auth/config/AuthConfigs.kt

@@ -36,29 +36,26 @@ class AuthModeFactory(
   val airbyteEdition: Configs.AirbyteEdition,
 ) {
   /**
-   * When the Micronaut environment is set to `community-auth`, the `SIMPLE` auth mode is used
-   * regardless of the deployment mode or other configurations. This bean replaces the
-   * [defaultAuthMode] when the `community-auth` environment is active.
+   * When the Airbyte edition is set to `community` and Micronaut Security is enabled, the
+   * `SIMPLE` auth mode is used regardless of the deployment mode or other configurations.
    */
   @Singleton
-  @Requires(env = ["community-auth"])
+  @Requires(property = "airbyte.edition", value = "community")
+  @Requires(property = "micronaut.security.enabled", value = "true")
   @Primary
-  fun communityAuthMode(): AuthMode {
-    return AuthMode.SIMPLE
-  }
+  fun communityAuthMode(): AuthMode = AuthMode.SIMPLE
 
   /**
    * The default auth mode is determined by the deployment mode and edition.
    */
   @Singleton
-  fun defaultAuthMode(): AuthMode {
-    return when {
+  fun defaultAuthMode(): AuthMode =
+    when {
       deploymentMode == DeploymentMode.CLOUD -> AuthMode.OIDC
       airbyteEdition == Configs.AirbyteEdition.PRO -> AuthMode.OIDC
       deploymentMode == DeploymentMode.OSS -> AuthMode.NONE
       else -> throw IllegalStateException("Unknown or unspecified deployment mode: $deploymentMode")
     }
-  }
 }
 
 /**
@@ -75,7 +72,5 @@ class AuthConfigFactory(
   val initialUserConfig: InitialUserConfig? = null,
 ) {
   @Singleton
-  fun authConfig(): AuthConfigs {
-    return AuthConfigs(authMode, keycloakConfig, oidcConfig, initialUserConfig)
-  }
+  fun authConfig(): AuthConfigs = AuthConfigs(authMode, keycloakConfig, oidcConfig, initialUserConfig)
 }

airbyte-commons-auth/src/test/kotlin/io/airbyte/commons/auth/config/AuthConfigsTest.kt

@@ -86,13 +86,15 @@ class AuthConfigsForCloudTest {
   }
 }
 
-@MicronautTest(environments = ["community-auth"])
+@MicronautTest
 class AuthConfigsForCommunityAuthTest {
   @Inject
   lateinit var authConfigs: AuthConfigs
 
   @Test
   @Property(name = "airbyte.deployment-mode", value = "OSS")
+  @Property(name = "airbyte.edition", value = "community")
+  @Property(name = "micronaut.security.enabled", value = "true")
   fun `test community-auth environment sets mode to SIMPLE`() {
     Assertions.assertTrue(authConfigs.authMode == AuthMode.SIMPLE)
   }

airbyte-commons-micronaut/build.gradle.kts

@@ -5,7 +5,7 @@ plugins {
 
 dependencies {
   compileOnly(libs.lombok)
-  annotationProcessor(libs.lombok)     // Lombok must be added BEFORE Micronaut
+  annotationProcessor(libs.lombok) // Lombok must be added BEFORE Micronaut
   annotationProcessor(platform(libs.micronaut.platform))
   annotationProcessor(libs.bundles.micronaut.annotation.processor)
 
@@ -25,6 +25,7 @@ dependencies {
 
   testImplementation(libs.bundles.micronaut.test)
   testImplementation(libs.mockito.inline)
+  testImplementation(libs.mockk)
 }
 
 tasks.named<Test>("test") {

airbyte-commons-micronaut/src/main/kotlin/io/airbyte/micronaut/config/EditionPropertySourceLoader.kt

@@ -0,0 +1,36 @@
+package io.airbyte.micronaut.config
+
+import io.micronaut.context.env.AbstractPropertySourceLoader
+import io.micronaut.context.env.PropertySource
+import io.micronaut.context.env.yaml.YamlPropertySourceLoader
+import io.micronaut.core.io.ResourceLoader
+import java.util.Optional
+
+/**
+ * Loads properties based on the edition of Airbyte.
+ * Looks for a file named `application-edition-<edition>.yml` where `<edition>` is the value of the
+ * `AIRBYTE_EDITION` environment variable.
+ * This loader is registered in META-INF/services/io.micronaut.context.env.PropertySourceLoader
+ * so that it is automatically picked up by Micronaut.
+ */
+class EditionPropertySourceLoader(
+  private val airbyteEdition: String? = System.getenv(AIRBYTE_EDITION_ENV_VAR),
+) : YamlPropertySourceLoader() {
+  companion object {
+    const val AIRBYTE_EDITION_ENV_VAR = "AIRBYTE_EDITION"
+  }
+
+  override fun getOrder(): Int = AbstractPropertySourceLoader.DEFAULT_POSITION + 10
+
+  override fun load(
+    name: String?,
+    resourceLoader: ResourceLoader?,
+  ): Optional<PropertySource> {
+    if (airbyteEdition.isNullOrEmpty()) {
+      return Optional.empty()
+    }
+
+    val fileName = "application-edition-" + airbyteEdition.lowercase()
+    return super.load(fileName, resourceLoader)
+  }
+}

airbyte-commons-micronaut/src/main/resources/META-INF/services/io.micronaut.context.env.PropertySourceLoader

@@ -1 +1,2 @@
-io.airbyte.micronaut.config.AirbytePropertySourceLoader
+io.airbyte.micronaut.config.AirbytePropertySourceLoader
+io.airbyte.micronaut.config.EditionPropertySourceLoader

airbyte-commons-micronaut/src/test/kotlin/io/airbyte/micronaut/config/EditionPropertySourceLoaderTest.kt

@@ -0,0 +1,43 @@
+package io.airbyte.micronaut.config
+
+import io.micronaut.core.io.ResourceLoader
+import io.mockk.Called
+import io.mockk.mockk
+import io.mockk.verify
+import org.junit.jupiter.api.Assertions.assertTrue
+import org.junit.jupiter.api.Test
+
+class EditionPropertySourceLoaderTest {
+  private val resourceLoader = mockk<ResourceLoader>(relaxed = true)
+
+  @Test
+  fun `should return empty optional when AIRBYTE_EDITION is not set`() {
+    val loader = EditionPropertySourceLoader(null)
+
+    val result = loader.load("test", resourceLoader)
+
+    assertTrue(result.isEmpty)
+    verify { resourceLoader wasNot Called }
+  }
+
+  @Test
+  fun `should return empty optional when AIRBYTE_EDITION is empty`() {
+    val loader = EditionPropertySourceLoader("")
+
+    val result = loader.load("test", resourceLoader)
+
+    assertTrue(result.isEmpty)
+    verify { resourceLoader wasNot Called }
+  }
+
+  @Test
+  fun `should attempt to load correct file when AIRBYTE_EDITION is set`() {
+    val loader = EditionPropertySourceLoader("community")
+
+    loader.load("test", resourceLoader)
+
+    // ensures underlying loader is called with correct file name and extensions
+    verify { resourceLoader.getResourceAsStream("application-edition-community.yml") }
+    verify { resourceLoader.getResourceAsStream("application-edition-community.yaml") }
+  }
+}

airbyte-commons-server/build.gradle.kts

@@ -25,6 +25,7 @@ dependencies {
   implementation(libs.micronaut.inject)
   implementation(libs.micronaut.jaxrs.server)
   implementation(libs.micronaut.security)
+  implementation(libs.micronaut.security.jwt)
   implementation(libs.bundles.micronaut.data.jdbc)
   implementation(libs.bundles.micronaut.kotlin)
   implementation(libs.bundles.flyway)

airbyte-server/build.gradle.kts

@@ -9,7 +9,7 @@ plugins {
 
 dependencies {
   compileOnly(libs.lombok)
-  annotationProcessor(libs.lombok)     // Lombok must be added BEFORE Micronaut
+  annotationProcessor(libs.lombok) // Lombok must be added BEFORE Micronaut
   annotationProcessor(platform(libs.micronaut.platform))
   annotationProcessor(libs.bundles.micronaut.annotation.processor)
   annotationProcessor(libs.micronaut.jaxrs.processor)
@@ -26,6 +26,7 @@ dependencies {
   implementation(libs.micronaut.jaxrs.server)
   implementation(libs.micronaut.http)
   implementation(libs.micronaut.security)
+  implementation(libs.micronaut.security.jwt)
   implementation(libs.bundles.flyway)
   implementation(libs.s3)
   implementation(libs.sts)
@@ -78,7 +79,7 @@ dependencies {
   runtimeOnly(libs.h2.database)
 
   testCompileOnly(libs.lombok)
-  testAnnotationProcessor(libs.lombok)     // Lombok must be added BEFORE Micronaut
+  testAnnotationProcessor(libs.lombok) // Lombok must be added BEFORE Micronaut
   testAnnotationProcessor(platform(libs.micronaut.platform))
   testAnnotationProcessor(libs.bundles.micronaut.annotation.processor)
   testAnnotationProcessor(libs.micronaut.jaxrs.processor)
@@ -116,9 +117,10 @@ tasks.named("assemble") {
   dependsOn(copySeed)
 }
 
-val env = Properties().apply {
-  load(rootProject.file(".env.dev").inputStream())
-}
+val env =
+  Properties().apply {
+    load(rootProject.file(".env.dev").inputStream())
+  }
 
 airbyte {
   application {
@@ -143,7 +145,7 @@ airbyte {
         "TRACKING_STRATEGY" to env["TRACKING_STRATEGY"].toString(),
         "TEMPORAL_HOST" to "localhost:7233",
         "MICRONAUT_ENVIRONMENTS" to "control-plane",
-      )
+      ),
     )
   }
 
@@ -157,11 +159,14 @@ airbyte {
   }
 
   spotbugs {
-      excludes = listOf("  <Match>\n" +
-              "    <Package name=\"io.airbyte.server.repositories.domain.*\" />\n" +
-              "    <!-- All args constructor used by builders trigger this error -->\n" +
-              "    <Bug pattern=\"NP_PARAMETER_MUST_BE_NONNULL_BUT_MARKED_AS_NULLABLE\" />\n" +
-              "  </Match>")
+    excludes =
+      listOf(
+        "  <Match>\n" +
+          "    <Package name=\"io.airbyte.server.repositories.domain.*\" />\n" +
+          "    <!-- All args constructor used by builders trigger this error -->\n" +
+          "    <Bug pattern=\"NP_PARAMETER_MUST_BE_NONNULL_BUT_MARKED_AS_NULLABLE\" />\n" +
+          "  </Match>",
+      )
   }
 }
 
@@ -171,7 +176,7 @@ tasks.named<Test>("test") {
       "AIRBYTE_VERSION" to env["VERSION"],
       "MICRONAUT_ENVIRONMENTS" to "test",
       "SERVICE_NAME" to project.name,
-    )
+    ),
   )
 }
 

airbyte-server/src/main/java/io/airbyte/server/pro/AirbyteAuthInternalTokenReader.java

@@ -5,7 +5,7 @@
 package io.airbyte.server.pro;
 
 import io.airbyte.commons.auth.AirbyteAuthConstants;
-import io.airbyte.commons.license.annotation.RequiresAirbyteProEnabled;
+import io.micronaut.context.annotation.Requires;
 import io.micronaut.security.token.reader.HttpHeaderTokenReader;
 import jakarta.inject.Singleton;
 
@@ -17,7 +17,10 @@
  * will only be present on internal requests.
  */
 @Singleton
-@RequiresAirbyteProEnabled
+@Requires(property = "micronaut.security.enabled",
+          value = "true")
+@Requires(property = "airbyte.deployment-mode",
+          value = "OSS")
 public class AirbyteAuthInternalTokenReader extends HttpHeaderTokenReader {
 
   // This is set higher than other token readers so that it is checked last.

airbyte-server/src/main/java/io/airbyte/server/pro/AirbyteAuthInternalTokenValidator.java

@@ -7,8 +7,8 @@
 import static io.airbyte.config.persistence.UserPersistence.DEFAULT_USER_ID;
 
 import io.airbyte.commons.auth.AirbyteAuthConstants;
-import io.airbyte.commons.license.annotation.RequiresAirbyteProEnabled;
 import io.airbyte.commons.server.support.RbacRoleHelper;
+import io.micronaut.context.annotation.Requires;
 import io.micronaut.http.HttpRequest;
 import io.micronaut.security.authentication.Authentication;
 import io.micronaut.security.token.validator.TokenValidator;
@@ -26,7 +26,10 @@
  **/
 @Slf4j
 @Singleton
-@RequiresAirbyteProEnabled
+@Requires(property = "micronaut.security.enabled",
+          value = "true")
+@Requires(property = "airbyte.deployment-mode",
+          value = "OSS")
 public class AirbyteAuthInternalTokenValidator implements TokenValidator<HttpRequest<?>> {
 
   @Override

airbyte-server/src/main/kotlin/io/airbyte/server/config/community/auth/CommunityAuthLogoutEventListener.kt

@@ -0,0 +1,25 @@
+package io.airbyte.server.config.community.auth
+
+import io.airbyte.api.problems.model.generated.ProblemMessageData
+import io.airbyte.api.problems.throwable.generated.UnprocessableEntityProblem
+import io.airbyte.data.services.AuthRefreshTokenService
+import io.micronaut.context.event.ApplicationEventListener
+import io.micronaut.security.authentication.Authentication
+import io.micronaut.security.event.LogoutEvent
+import jakarta.inject.Singleton
+
+@Singleton
+class CommunityAuthLogoutEventListener(
+  private val refreshTokenService: AuthRefreshTokenService,
+) : ApplicationEventListener<LogoutEvent> {
+  override fun onApplicationEvent(event: LogoutEvent) {
+    val sessionId =
+      (event.source as? Authentication)
+        ?.attributes
+        ?.get(SESSION_ID)
+        ?.toString()
+        ?: throw UnprocessableEntityProblem(ProblemMessageData().message("Could not retrieve session ID from authentication context"))
+
+    refreshTokenService.revokeAuthRefreshToken(sessionId)
+  }
+}

airbyte-server/src/main/kotlin/io/airbyte/server/SimpleAuthProvider.kt → airbyte-server/src/main/kotlin/io/airbyte/server/config/community/auth/CommunityAuthProvider.kt

@@ -1,7 +1,7 @@
 /*
  * Copyright (c) 2020-2024 Airbyte, Inc., all rights reserved.
  */
-package io.airbyte.server
+package io.airbyte.server.config.community.auth
 
 import io.airbyte.commons.auth.RequiresAuthMode
 import io.airbyte.commons.auth.config.AuthMode
@@ -13,22 +13,31 @@ import io.micronaut.security.authentication.AuthenticationRequest
 import io.micronaut.security.authentication.AuthenticationResponse
 import io.micronaut.security.authentication.provider.HttpRequestAuthenticationProvider
 import jakarta.inject.Singleton
+import java.util.UUID
+
+const val SESSION_ID = "sessionId"
 
 /**
  * This class is responsible for authenticating the user against the community provider.
  */
 @Singleton
 @RequiresAuthMode(AuthMode.SIMPLE)
-class SimpleAuthProvider<B>(private val instanceAdminConfig: InstanceAdminConfig) : HttpRequestAuthenticationProvider<B> {
+class CommunityAuthProvider<B>(
+  private val instanceAdminConfig: InstanceAdminConfig,
+) : HttpRequestAuthenticationProvider<B> {
   override fun authenticate(
     requestContext: HttpRequest<B>?,
     authRequest: AuthenticationRequest<String, String>,
   ): AuthenticationResponse? {
     if (authRequest.identity == instanceAdminConfig.username && authRequest.secret == instanceAdminConfig.password) {
-      return AuthenticationResponse.success(
-        UserPersistence.DEFAULT_USER_ID.toString(),
-        RbacRoleHelper.getInstanceAdminRoles(),
-      )
+      val sessionId = UUID.randomUUID()
+      val authenticationResponse =
+        AuthenticationResponse.success(
+          UserPersistence.DEFAULT_USER_ID.toString(),
+          RbacRoleHelper.getInstanceAdminRoles(),
+          mapOf(SESSION_ID to sessionId.toString()),
+        )
+      return authenticationResponse
     }
     return AuthenticationResponse.failure("Invalid credentials")
   }