Role-based access control with Caliban

rbac.sc

//> using dep com.github.ghostdogpr::caliban-quick:2.7.1

import caliban.*
import caliban.CalibanError.*
import caliban.Value.StringValue
import caliban.execution.FieldInfo
import caliban.parsing.adt.Directive
import caliban.quick.*
import caliban.schema.Annotations.*
import caliban.schema.Schema
import caliban.wrappers.Wrapper.FieldWrapper
import scala.util.Try
import zio.*
import zio.http.*
import zio.query.ZQuery

enum Role {
  case Admin, User
}

val directiveName = "hasRole"
val attributeName = "role"

class HasRoleDirective(role: Role)
    extends GQLDirective(Directive(directiveName, Map(attributeName -> StringValue(role.toString))))

case class admin() extends HasRoleDirective(Role.Admin)
case class user()  extends HasRoleDirective(Role.User)

case class AuthContext(roles: Set[Role])

def getRequiredRoles(info: FieldInfo): Set[Role] =
  info.directives
    .filter(_.name == directiveName)
    .flatMap(_.arguments.get(attributeName))
    .flatMap {
      case StringValue(role) => Try(Role.valueOf(role)).toOption.toList
      case _                 => Nil
    }
    .toSet

val accessControl: FieldWrapper[AuthContext] =
  new FieldWrapper[AuthContext](wrapPureValues = true) {
    def wrap[R1 <: AuthContext](
        query: ZQuery[R1, ExecutionError, ResponseValue],
        info: FieldInfo
    ): ZQuery[R1, ExecutionError, ResponseValue] =
      ZQuery.serviceWithQuery[AuthContext] { ctx =>
        val missingRoles = getRequiredRoles(info).diff(ctx.roles)
        if (missingRoles.isEmpty) query
        else ZQuery.fail(ExecutionError(s"Missing required roles: ${missingRoles.mkString(", ")}."))
      }
  }

case class Query(
    @admin adminData: String,
    @user userData: String
) derives Schema.SemiAuto

val api = graphQL(RootResolver(Query("admin", "user"))) @@ accessControl

val middleware =
  Middleware.customAuthProviding[AuthContext] { req =>
    req.headers
      .get("Roles")
      .map(_.split(",").flatMap(s => Try(Role.valueOf(s)).toOption).toSet)
      .map(roles => AuthContext(roles))
  }

Unsafe.unsafely {
  Runtime.default.unsafe.run {
    api
      .routes("api/graphql")
      .map(_ @@ middleware)
      .flatMap(_.serve.provideLayer(Server.defaultWithPort(8080)))
  }
}