feat: azure rbac aks service discovery (#1311)

.gitignore

@@ -21,3 +21,7 @@ docs/src/main/paradox/attachments
 .metals
 metals.sbt
 **/project/project
+
+
+# Bloop
+.bloop

build.sbt

@@ -1,5 +1,8 @@
-import com.typesafe.sbt.packager.docker.{ Cmd, ExecCmd }
+import com.typesafe.sbt.packager.docker.{Cmd, ExecCmd}
 import sbt.Keys.parallelExecution
+import com.typesafe.tools.mima.core.ProblemFilter
+import com.typesafe.tools.mima.plugin.MimaPlugin
+import com.typesafe.tools.mima.plugin.MimaPlugin.autoImport.*
 
 ThisBuild / resolvers += "Akka library repository".at("https://repo.akka.io/maven")
 ThisBuild / resolvers += Resolver.jcenterRepo
@@ -19,6 +22,7 @@ lazy val `akka-management-root` = project
     // in AkkaManagement should also be updated
     `akka-discovery-aws-api`,
     `akka-discovery-aws-api-async`,
+    `akka-discovery-azure-api`,
     `akka-discovery-kubernetes-api`,
     `akka-discovery-marathon-api`,
     `akka-management`,
@@ -71,6 +75,18 @@ lazy val `akka-discovery-kubernetes-api` = project
   )
   .dependsOn(`akka-management-pki`)
 
+lazy val `akka-discovery-azure-api` = (project in file("discovery-azure-api"))
+  .enablePlugins(AutomateHeaderPlugin)
+  .disablePlugins(com.geirsson.CiReleasePlugin)
+  .settings(
+    name := "akka-discovery-azure-api",
+    organization := "com.lightbend.akka.discovery",
+    libraryDependencies := Dependencies.DiscoveryAzureApi,
+        // FIXME: update once we have a release out
+        mimaPreviousArtifacts := Set.empty
+  )
+  .dependsOn(`akka-management-pki`)
+
 lazy val `akka-discovery-marathon-api` = project
   .in(file("discovery-marathon-api"))
   .enablePlugins(AutomateHeaderPlugin)

discovery-azure-api/src/main/resources/reference.conf

@@ -0,0 +1,64 @@
+######################################################
+# Akka Service Discovery Azure Config
+######################################################
+
+akka.discovery {
+  # Set the following in your application.conf if you want to use this discovery mechanism:
+  # method = azure-rbac-aks-api
+  azure-rbac-aks-api {
+    class = akka.discovery.azureapi.rbac.aks.AzureRbacAksServiceDiscovery
+
+    authority-host = "https://login.microsoftonline.com/"
+    authority-host = ${?AZURE_AUTHORITY_HOST}
+
+    # Required
+    # Injected by the workload identity controller manager
+    client-id = ${AZURE_CLIENT_ID}
+
+    federated-token-file = "/var/run/secrets/azure/tokens/azure-identity-token"
+    federated-token-file = ${?AZURE_FEDERATED_TOKEN_FILE}
+
+    # Required
+    # Injected by the workload identity controller manager
+    tenant-id = ${AZURE_TENANT_ID}
+
+    # AKS uses a pair of first-party Microsoft Entra applications
+    # The AKS Microsoft Entra server application ID(scope) that the server side uses is 6dae42f8-4368-4678-94ff-3960e28e3630/.default
+    entra-server-id = "6dae42f8-4368-4678-94ff-3960e28e3630/.default"
+    entra-server-id = ${?AZURE_SERVER_ID}
+
+    # API server, cert and token information. Currently these are present on K8s versions: 1.6, 1.7, 1.8, and perhaps more
+    api-ca-path = "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"
+    api-token-path = "/var/run/secrets/kubernetes.io/serviceaccount/token"
+
+    # Required
+    api-service-host = ${KUBERNETES_SERVICE_HOST}
+    api-service-port = ${KUBERNETES_SERVICE_PORT}
+
+    # Namespace discovery path
+    #
+    # If this path doesn't exist, the namespace will default to "default".
+    pod-namespace-path = "/var/run/secrets/kubernetes.io/serviceaccount/namespace"
+
+    # Namespace to query for pods.
+    #
+    # Set this value to a specific string to override discovering the namespace using pod-namespace-path.
+    pod-namespace = "<pod-namespace>"
+    pod-namespace = ${?KUBERNETES_NAMESPACE}
+
+    # Domain of the k8s cluster
+    pod-domain = "cluster.local"
+
+    # Selector value to query pod API with.
+    # `%s` will be replaced with the configured effective name, which defaults to the actor system name
+    pod-label-selector = "app=%s"
+
+    # Enables the usage of the raw IP instead of the composed value for the resolved target host
+    # Note that when using names, the deprecated DNS form <a>-<b>-<c>-<d>.<ns>.pod.<zone> is used
+    # and that may not work on newer Kubernetes versions.
+    use-raw-ip = true
+
+    # When set, validate the container is not in 'waiting' state
+    container-name = ""
+  }
+}

discovery-azure-api/src/main/scala/akka/discovery/azureapi/rbac/aks/AzureRbacAksServiceDiscovery.scala

@@ -0,0 +1,254 @@
+/*
+ * Copyright (C) 2017-2024 Lightbend Inc. <https://www.lightbend.com>
+ */
+
+package akka.discovery.azureapi.rbac.aks
+
+import akka.actor.ExtendedActorSystem
+import akka.annotation.InternalApi
+import akka.discovery.ServiceDiscovery.{ Resolved, ResolvedTarget }
+import akka.discovery.azureapi.rbac.aks.AzureRbacAksServiceDiscovery._
+import akka.discovery.azureapi.rbac.aks.JsonFormat._
+import akka.discovery.{ Lookup, ServiceDiscovery }
+import akka.dispatch.Dispatchers.DefaultBlockingDispatcherId
+import akka.event.Logging
+import akka.http.scaladsl.model._
+import akka.http.scaladsl.model.headers.{ Authorization, OAuth2BearerToken }
+import akka.http.scaladsl.unmarshalling.Unmarshal
+import akka.http.scaladsl.{ ConnectionContext, Http, HttpsConnectionContext }
+import akka.pki.kubernetes.PemManagersProvider
+import com.azure.core.credential.{ AccessToken, TokenRequestContext }
+import com.azure.identity.{ DefaultAzureCredential, DefaultAzureCredentialBuilder }
+
+import java.net.InetAddress
+import java.nio.file.{ Files, Paths }
+import java.security.{ KeyStore, SecureRandom }
+import javax.net.ssl.{ KeyManager, KeyManagerFactory, SSLContext, TrustManager }
+import scala.collection.immutable
+import scala.concurrent.duration.FiniteDuration
+import scala.concurrent.{ ExecutionContext, Future }
+import scala.jdk.FutureConverters._
+import scala.util.control.{ NoStackTrace, NonFatal }
+
+/**
+ * INTERNAL API
+ *
+ * Finds relevant targets given a pod list. Note that this doesn't filter by name as it is the job of the selector
+ * to do that.
+ */
+@InternalApi
+object AzureRbacAksServiceDiscovery {
+  private def azureDefaultCredential: DefaultAzureCredential =
+    new DefaultAzureCredentialBuilder().build()
+
+  private val accessTokenRequestContext: TokenRequestContext =
+    new TokenRequestContext()
+
+  private[aks] def targets(
+      podList: PodList,
+      portName: Option[String],
+      podNamespace: String,
+      podDomain: String,
+      rawIp: Boolean,
+      containerName: Option[String]): immutable.Seq[ResolvedTarget] =
+    for {
+      item <- podList.items
+      if item.metadata.flatMap(_.deletionTimestamp).isEmpty
+      itemSpec <- item.spec.toSeq
+      itemStatus <- item.status.toSeq
+      if itemStatus.phase.contains("Running")
+      if containerName.forall(name =>
+        itemStatus.containerStatuses match {
+          case Some(statuses) => statuses.filter(_.name == name).exists(!_.state.contains("waiting"))
+          case None           => false
+        })
+      ip <- itemStatus.podIP.toSeq
+      // Maybe port is an Option of a port, and will be None if no portName was requested
+      maybePort <- portName match {
+        case None =>
+          List(None)
+        case Some(name) =>
+          for {
+            container <- itemSpec.containers
+            ports <- container.ports.toSeq
+            port <- ports
+            if port.name.contains(name)
+          } yield Some(port.containerPort)
+      }
+    } yield {
+      val hostOrIp = if (rawIp) ip else s"${ip.replace('.', '-')}.$podNamespace.pod.$podDomain"
+      ResolvedTarget(
+        host = hostOrIp,
+        port = maybePort,
+        address = Some(InetAddress.getByName(ip))
+      )
+    }
+
+  private final class KubernetesApiException(msg: String) extends RuntimeException(msg) with NoStackTrace
+
+  private final class AzureIdentityException(msg: String) extends RuntimeException(msg) with NoStackTrace
+
+  private final case class KubernetesSetup(namespace: String, ctx: HttpsConnectionContext)
+}
+
+/**
+ * INTERNAL API
+ *
+ * Finds relevant targets given a pod list. Note that this doesn't filter by name as it is the job of the selector
+ * to do that.
+ */
+@InternalApi
+final class AzureRbacAksServiceDiscovery(implicit system: ExtendedActorSystem) extends ServiceDiscovery {
+  private val http = Http()
+
+  private val settings = Settings(system.settings.config.getConfig("akka.discovery.azure-rbac-aks-api"))
+
+  private val log = Logging(system, classOf[AzureRbacAksServiceDiscovery])
+
+  private val kubernetesSetup = {
+    implicit val blockingDispatcher: ExecutionContext = system.dispatchers.lookup(DefaultBlockingDispatcherId)
+    for {
+      namespace: String <- Future {
+        settings.podNamespace
+          .orElse(readConfigVarFromFilesystem(settings.podNamespacePath, "pod-namespace"))
+          .getOrElse("default")
+      }
+      httpsContext <- Future(clientHttpsConnectionContext())
+    } yield {
+      KubernetesSetup(namespace, httpsContext)
+    }
+  }
+
+  log.debug("Settings [{}]", settings)
+
+  import system.dispatcher
+
+  private def fetchAccessToken: Future[AccessToken] =
+    azureDefaultCredential
+      .getToken(accessTokenRequestContext.addScopes(settings.entraServerId))
+      .onErrorMap { error =>
+        log.error("[{}]", error)
+        new AzureIdentityException(
+          "Attempt failed while fetching access token. Check if workload identity is enabled for the cluster or not and if the pods has been injected with required AZURE environment variables"
+        )
+      }
+      .toFuture
+      .asScala
+
+  private def parseKubernetesResponse(response: HttpResponse, entity: HttpEntity.Strict): Future[PodList] =
+    response.status match {
+      case StatusCodes.OK =>
+        log.debug("Kubernetes API entity: [{}]", entity.data.utf8String)
+        Unmarshal(entity).to[PodList].recoverWith {
+          case exception =>
+            log.warning(
+              "Failed to unmarshal Kubernetes API response.  Status code: [{}]; Response body: [{}]. Ex: [{}]",
+              response.status.value,
+              entity,
+              exception.getMessage)
+
+            Future.failed(new KubernetesApiException("Failed to unmarshal Kubernetes API response."))
+        }
+      case StatusCodes.Forbidden =>
+        Unmarshal(entity).to[String].flatMap { body =>
+          log.warning("Forbidden to communicate with Kubernetes API server; check RBAC settings. Response: [{}]", body)
+
+          Future.failed(
+            new KubernetesApiException(
+              "Forbidden when communicating with the Kubernetes API Server. Check if the managed identity has the appropriate role assigment(example: Azure Pod Reader) or if workload identity is enabled for the cluster."
+            )
+          )
+        }
+      case other =>
+        Unmarshal(entity).to[String].flatMap { body =>
+          log.warning(
+            "Non-200 when communicating with Kubernetes API server. Status code: [{}]. Response body: [{}]",
+            other,
+            body
+          )
+
+          Future.failed(new KubernetesApiException(s"Non-200 from Kubernetes API server: $other"))
+        }
+    }
+
+  private def pods(ctx: HttpsConnectionContext, request: HttpRequest, timeout: FiniteDuration): Future[PodList] = {
+    for {
+      response: HttpResponse <- http.singleRequest(request, ctx)
+      entity <- response.entity.toStrict(timeout)
+      pods <- parseKubernetesResponse(response, entity)
+    } yield pods
+  }
+
+  override def lookup(lookup: Lookup, resolveTimeout: FiniteDuration): Future[ServiceDiscovery.Resolved] = {
+    val selector = settings.podLabelSelector.format(lookup.serviceName)
+
+    for {
+      ks <- kubernetesSetup
+      token <- fetchAccessToken.map(_.getToken)
+      request <- podRequest(token, ks.namespace, selector)
+      pods <- pods(ks.ctx, request, resolveTimeout)
+    } yield {
+      val addresses =
+        targets(pods, lookup.portName, ks.namespace, settings.podDomain, settings.rawIp, settings.containerName)
+      if (addresses.isEmpty && pods.items.nonEmpty) {
+        val containerPortNames = pods.items.flatMap(_.spec).flatMap(_.containers).flatMap(_.ports).flatten.toSet
+        log.warning(
+          "No targets found from pod list. Is the correct port name configured? Current configuration: [{}]. Ports on pods: [{}]",
+          lookup.portName,
+          containerPortNames
+        )
+      }
+      Resolved(
+        serviceName = lookup.serviceName,
+        addresses = addresses
+      )
+    }
+  }
+
+  private def podRequest(token: String, namespace: String, labelSelector: String) = {
+    val host = settings.apiServiceHost
+    val port = settings.apiServicePort
+    val path = Uri.Path.Empty / "api" / "v1" / "namespaces" / namespace / "pods"
+    val query = Uri.Query("labelSelector" -> labelSelector)
+    val uri = Uri.from(scheme = "https", host = host, port = port).withPath(path).withQuery(query)
+
+    Future(HttpRequest(uri = uri, headers = List(Authorization(OAuth2BearerToken(token)))))
+  }
+
+  /**
+   * This uses blocking IO, and so should only be used at startup from blocking dispatcher.
+   */
+  private def clientHttpsConnectionContext(): HttpsConnectionContext = {
+    val certificates = PemManagersProvider.loadCertificates(settings.apiCaPath)
+    val factory = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm)
+    val keyStore = KeyStore.getInstance("PKCS12")
+    keyStore.load(null)
+    factory.init(keyStore, Array.empty)
+    val km: Array[KeyManager] = factory.getKeyManagers
+    val tm: Array[TrustManager] =
+      PemManagersProvider.buildTrustManagers(certificates)
+    val random: SecureRandom = new SecureRandom
+    val sslContext = SSLContext.getInstance("TLSv1.2")
+    sslContext.init(km, tm, random)
+    ConnectionContext.httpsClient(sslContext)
+  }
+
+  /**
+   * This uses blocking IO, and so should only be used to read configuration at startup from blocking dispatcher.
+   */
+  private def readConfigVarFromFilesystem(path: String, name: String): Option[String] = {
+    val file = Paths.get(path)
+    if (Files.exists(file)) {
+      try {
+        Some(Files.readString(file))
+      } catch {
+        case NonFatal(e) =>
+          log.error(e, "Error reading [{}] from [{}]", name, path)
+          None
+      }
+    } else {
+      log.warning("Unable to read [{}] from [{}] because it doesn't exist.", name, path)
+      None
+    }
+  }
+}

discovery-azure-api/src/main/scala/akka/discovery/azureapi/rbac/aks/JsonFormat.scala

@@ -0,0 +1,24 @@
+/*
+ * Copyright (C) 2017-2024 Lightbend Inc. <https://www.lightbend.com>
+ */
+
+package akka.discovery.azureapi.rbac.aks
+
+import akka.annotation.InternalApi
+import PodList._
+import akka.http.scaladsl.marshallers.sprayjson.SprayJsonSupport
+import spray.json._
+
+/**
+ * INTERNAL API
+ */
+@InternalApi private[aks] object JsonFormat extends SprayJsonSupport with DefaultJsonProtocol {
+  implicit val containerPortFormat: JsonFormat[ContainerPort] = jsonFormat2(ContainerPort.apply)
+  implicit val containerFormat: JsonFormat[Container] = jsonFormat2(Container.apply)
+  implicit val podSpecFormat: JsonFormat[PodSpec] = jsonFormat1(PodSpec.apply)
+  implicit val containerStatusFormat: JsonFormat[ContainerStatus] = jsonFormat2(ContainerStatus.apply)
+  implicit val podStatusFormat: JsonFormat[PodStatus] = jsonFormat3(PodStatus.apply)
+  implicit val metadataFormat: JsonFormat[Metadata] = jsonFormat2(Metadata.apply)
+  implicit val podFormat: JsonFormat[Pod] = jsonFormat3(Pod.apply)
+  implicit val podListFormat: RootJsonFormat[PodList] = jsonFormat1(PodList.apply)
+}