feat: auth for workload APIs (#16702)

Co-authored-by: Cam Kennedy <[email protected]>

Author: abuchanan-airbyte

airbyte-api/workload-api/src/main/openapi/workload-openapi.yaml

@@ -248,7 +248,7 @@ paths:
               $ref: "#/components/schemas/WorkloadQueueCleanLimit"
       responses:
         "204":
-          description: Cleaning workload queue successfull
+          description: Cleaning workload queue successful
   /api/v1/workload/queue/depth:
     post:
       tags:

airbyte-commons-server/src/main/java/io/airbyte/commons/server/authorization/KeycloakTokenValidator.java

@@ -54,8 +54,16 @@
 @Singleton
 @Primary
 @RequiresAuthMode(AuthMode.OIDC)
+// We're not confident about what the identity-provider.type will be when keycloak *should* be
+// enabled,
+// (we think it's usually "oidc" or "keycloak"). We're more confident about when we definitely
+// *don't* want it enabled,
+// so here we rule out "generic-oidc" and "simple" explicitly. Otherwise, for now, keycloak is
+// enabled.
 @Requires(property = "airbyte.auth.identity-provider.type",
           notEquals = "generic-oidc")
+@Requires(property = "airbyte.auth.identity-provider.type",
+          notEquals = "simple")
 @SuppressWarnings({"PMD.PreserveStackTrace", "PMD.UseTryWithResources", "PMD.UnusedFormalParameter", "PMD.UnusedPrivateMethod",
   "PMD.ExceptionAsFlowControl"})
 public class KeycloakTokenValidator implements TokenValidator<HttpRequest<?>> {

airbyte-commons-server/src/main/kotlin/io/airbyte/commons/server/authorization/AuthenticationFactory.kt

@@ -29,21 +29,18 @@ class AuthenticationFactory(
   }
 
   private fun createAuth(
-    authUserId: String,
-    attrs: Map<String, Any>,
+    subject: String,
+    claims: Map<String, Any>,
   ): Authentication {
-    // Some tokens already have roles assigned. If the token contains roles, use those,
-    // otherwise resolve the roles for the current identity + request.
-    val tokenRoles = (attrs["roles"] as? List<*>)?.filterIsInstance<String>()
-    val roles = tokenRoles ?: resolveRoles(authUserId)
+    logger.debug { "Creating auth for $subject" }
 
-    return Authentication.build(authUserId, roles, attrs)
-  }
+    val roles =
+      roleResolver
+        .newRequest()
+        .withClaims(subject, claims)
+        .withRefsFromCurrentHttpRequest()
+        .roles()
 
-  private fun resolveRoles(authUserId: String): Set<String> =
-    roleResolver
-      .Request()
-      .withAuthUserId(authUserId)
-      .withCurrentHttpRequest()
-      .roles()
+    return Authentication.build(subject, roles, claims)
+  }
 }

airbyte-commons-server/src/main/kotlin/io/airbyte/commons/server/authorization/RoleResolver.kt

@@ -6,6 +6,7 @@ package io.airbyte.commons.server.authorization
 
 import io.airbyte.api.problems.model.generated.ProblemMessageData
 import io.airbyte.api.problems.throwable.generated.ForbiddenProblem
+import io.airbyte.commons.DEFAULT_USER_ID
 import io.airbyte.commons.auth.AuthRole
 import io.airbyte.commons.auth.AuthRoleConstants
 import io.airbyte.commons.auth.OrganizationAuthRole
@@ -19,9 +20,14 @@ import io.airbyte.commons.server.support.CurrentUserService
 import io.airbyte.config.Permission
 import io.airbyte.config.Permission.PermissionType
 import io.airbyte.config.helpers.PermissionHelper
+import io.airbyte.data.TokenType
+import io.airbyte.featureflag.FeatureFlagClient
+import io.airbyte.featureflag.IgnoreTokenRoleClaims
+import io.airbyte.featureflag.TokenSubject
 import io.github.oshai.kotlinlogging.KotlinLogging
 import io.micronaut.http.HttpRequest
 import io.micronaut.http.context.ServerRequestContext
+import io.micronaut.security.utils.SecurityService
 import jakarta.inject.Singleton
 import java.util.UUID
 
@@ -60,33 +66,81 @@ private val logger = KotlinLogging.logger {}
 open class RoleResolver(
   private val authenticationHeaderResolver: AuthenticationHeaderResolver,
   private val currentUserService: CurrentUserService,
+  private val securityService: SecurityService?,
   private val permissionHandler: PermissionHandler,
+  private val featureFlags: FeatureFlagClient,
 ) {
-  inner class Request {
-    var authUserId: String? = null
+  data class Subject(
+    val id: String,
+    val type: TokenType,
+  )
+
+  fun newRequest() = Request()
+
+  inner class Request internal constructor() {
+    var subject: Subject? = null
+    var claimedRoles: Set<String>? = null
     val props: MutableMap<String, String> = mutableMapOf()
+    val orgs: MutableSet<UUID> = mutableSetOf()
 
     fun withCurrentUser() =
       apply {
-        authUserId = currentUserService.currentUser.authUserId
+        subject = Subject(currentUserService.currentUser.authUserId, TokenType.USER)
       }
 
-    fun withAuthUserId(authUserId: String) =
+    fun withSubject(
+      id: String,
+      type: TokenType,
+    ) = apply {
+      this.subject = Subject(id, type)
+    }
+
+    fun withCurrentAuthentication() =
       apply {
-        this.authUserId = authUserId
+        // In community auth, where micronaut auth is disabled, the security service isn't available,
+        // so we need to manually fall back to the default user.
+        if (securityService == null) {
+          withSubject(DEFAULT_USER_ID.toString(), TokenType.USER)
+        }
+
+        securityService?.authentication?.map { auth ->
+          logger.debug { "Using current authentication object ${auth.name} ${auth.roles} ${auth.attributes}" }
+          withClaims(auth.name, auth.attributes)
+        }
       }
 
-    fun withCurrentHttpRequest() =
+    fun withClaims(
+      sub: String,
+      claims: Map<String, Any>,
+    ) = apply {
+      // Figure out the subject type and ID.
+      subject = Subject(sub, TokenType.fromClaims(claims))
+
+      // Some tokens have roles in the claims.
+      // If the token does have roles in the claims, then those are the only roles
+      // resolved by Request.roles().
+      val roles = (claims["roles"] as? List<*>)?.filterIsInstance<String>()
+      if (roles != null) {
+        claimedRoles = roles.toSet()
+      }
+    }
+
+    fun withRefsFromCurrentHttpRequest() =
       apply {
-        ServerRequestContext.currentRequest<Any>().map { withHttpRequest(it) }
+        ServerRequestContext.currentRequest<Any>().map { withRefsFromHttpRequest(it) }
       }
 
-    fun withHttpRequest(req: HttpRequest<*>) =
+    fun withRefsFromHttpRequest(req: HttpRequest<*>) =
       apply {
         val headers = req.headers.asMap(String::class.java, String::class.java)
         props.putAll(headers)
       }
 
+    fun withOrg(organizationId: UUID) =
+      apply {
+        orgs.add(organizationId)
+      }
+
     fun withRef(
       key: AuthenticationId,
       value: String,
@@ -110,34 +164,86 @@ open class RoleResolver(
      * roles() resolves the request details into a set of available roles.
      */
     fun roles(): Set<String> {
-      logger.debug { "Resolving roles for authUserId $authUserId" }
+      // these make null-checking cleaner
+      val subject = subject
+      var claimedRoles = claimedRoles
 
-      try {
-        val user = authUserId
-        if (user.isNullOrBlank()) {
-          logger.debug { "Provided authUserId is null or blank, returning empty role set" }
-          return emptySet()
-        }
+      // We'd like to transition away from tokens that hard-code roles in the claims,
+      // and look them up per-request instead. This will allow for more centralized,
+      // more consistent code for determining roles, and prevents bugs with having
+      // changing the set of roles needed by a token.
+      if (featureFlags.boolVariation(IgnoreTokenRoleClaims, TokenSubject(subject?.id ?: "unknown"))) {
+        claimedRoles = null
+      }
+
+      // This helps when new roles are added to the set of admin roles.
+      // "ADMIN" implies a set of other roles (see AuthRole.getInstanceAdminRoles()).
+      // When a new role is added to that set, any existing tokens in the wild that
+      // contain hard-coded roles will not have this new role. To deal with that,
+      // we regenerate the set of admin roles here.
+      //
+      // TODO Hopefully, in the near future, we'll move away from having hard-coded roles
+      //      in tokens and *always* generate the roles dynamically during the request processing,
+      //      so that this is no longer needed.
+      if (claimedRoles?.contains(AuthRoleConstants.ADMIN) == true) {
+        claimedRoles = AuthRole.getInstanceAdminRoles()
+      }
+
+      logger.debug { "Resolving roles for $subject" }
+
+      if (subject == null) {
+        logger.debug { "subject is null, returning empty role set" }
+        return emptySet()
+      }
+      if (subject.id.isBlank()) {
+        logger.debug { "subject.id is blank, returning empty role set" }
+        return emptySet()
+      }
+      if (!claimedRoles.isNullOrEmpty()) {
+        return claimedRoles
+      }
 
-        val workspaceIds = authenticationHeaderResolver.resolveWorkspace(props)?.toSet() ?: emptySet()
-        val organizationIds = authenticationHeaderResolver.resolveOrganization(props)?.toSet() ?: emptySet()
-        val authUserIds = authenticationHeaderResolver.resolveAuthUserIds(props) ?: emptySet()
-        val perms = permissionHandler.getPermissionsByAuthUserId(authUserId)
-        return resolveRoles(perms, user, workspaceIds, organizationIds, authUserIds)
+      return try {
+        when (subject.type) {
+          // Certain token types have a hard-coded list of roles.
+          TokenType.WORKLOAD_API -> setOf(AuthRoleConstants.DATAPLANE)
+          TokenType.EMBEDDED_V1 -> setOf(AuthRoleConstants.EMBEDDED_END_USER)
+          // Everything else resolves roles via the "permissions" table.
+          TokenType.DATAPLANE_V1, TokenType.SERVICE_ACCOUNT ->
+            resolvePermissions(
+              subject.id,
+              permissionHandler.getPermissionsByServiceAccountId(UUID.fromString(subject.id)),
+            )
+          TokenType.USER -> resolvePermissions(subject.id, permissionHandler.getPermissionsByAuthUserId(subject.id))
+        }
       } catch (e: Exception) {
-        logger.error(e) { "Failed to resolve roles for authUserId $authUserId" }
+        logger.error(e) { "Failed to resolve roles for $subject" }
         return emptySet()
       }
     }
 
+    private fun resolvePermissions(
+      subjectId: String,
+      perms: List<Permission>,
+    ): Set<String> {
+      logger.debug { "Resolving permissions for $subject and $perms" }
+
+      val workspaceIds = authenticationHeaderResolver.resolveWorkspace(props)?.toSet() ?: emptySet()
+      val resolvedOrgIds = authenticationHeaderResolver.resolveOrganization(props)?.toSet() ?: emptySet()
+      val authUserIds = authenticationHeaderResolver.resolveAuthUserIds(props) ?: emptySet()
+      val allOrgIds = orgs + resolvedOrgIds
+
+      return resolveRoles(perms, subjectId, workspaceIds, allOrgIds, authUserIds)
+    }
+
     /**
      * requireRole checks whether the request has the given role available,
      * and if not it throws ForbiddenProblem.
      */
     fun requireRole(role: String) {
       if (!roles().contains(role)) {
         throw ForbiddenProblem(
-          ProblemMessageData().message("User does not have the required $role permissions to access the resource(s)."),
+          ProblemMessageData().message("Caller does not have the required $role permissions to access the resource(s)."),
         )
       }
     }
@@ -155,33 +261,46 @@ open class RoleResolver(
    */
   fun resolveRoles(
     perms: List<Permission>,
-    currentAuthUserId: String,
+    subjectId: String,
     workspaceIds: Set<UUID>,
     organizationIds: Set<UUID>,
     authUserIds: Set<String>,
   ): Set<String> {
+    logger.debug {
+      "Resolving roles for $subjectId with perms=$perms workspaceIds=$workspaceIds organizationIds=$organizationIds authUserIds=$authUserIds"
+    }
+
     val roles = mutableSetOf(AuthRoleConstants.AUTHENTICATED_USER)
 
     // The SELF role denotes that request refers to the current request's identity.
     //
     // This relies on the assumption that the AuthenticationHeaderResolver will only
     // ever resolve authUserIds for one user (who can have multiple authUserIds).
     // TODO Technically, that's a weak assumption and we should make that interface clearer.
-    if (authUserIds.contains(currentAuthUserId)) {
+    if (authUserIds.contains(subjectId)) {
       roles.add(AuthRoleConstants.SELF)
     }
 
     if (perms.any { it.permissionType == PermissionType.INSTANCE_ADMIN }) {
       roles.addAll(AuthRole.getInstanceAdminRoles())
     }
 
+    perms.filter { it.permissionType == PermissionType.DATAPLANE }.forEach {
+      if (it.workspaceId == null && it.organizationId == null) {
+        roles.add(AuthRoleConstants.DATAPLANE)
+      }
+    }
+
     determineWorkspaceRole(perms, workspaceIds)?.let {
       roles.addAll(impliedRoles(it))
     }
     determineOrganizationRole(perms, organizationIds)?.let {
       roles.addAll(impliedRoles(it))
     }
 
+    logger.debug {
+      "Resolved roles for $subjectId with perms=$perms workspaceIds=$workspaceIds organizationIds=$organizationIds authUserIds=$authUserIds to $roles"
+    }
     return roles
   }
 }

airbyte-commons-server/src/main/kotlin/io/airbyte/commons/server/handlers/ResourceBootstrapHandler.kt

@@ -65,7 +65,7 @@ open class ResourceBootstrapHandler(
 
     // Ensure user has the required permissions to create a workspace
     roleResolver
-      .Request()
+      .newRequest()
       .withCurrentUser()
       .withRef(AuthenticationId.ORGANIZATION_ID, organization.organizationId)
       .requireRole(AuthRoleConstants.ORGANIZATION_ADMIN)