Cleaned up and working Azure token acquisition for engine

cloud-nio/cloud-nio-impl-drs/src/main/scala/cloud/nio/impl/drs/DrsCredentials.scala

@@ -2,13 +2,18 @@ package cloud.nio.impl.drs
 
 import cats.syntax.validated._
 import com.azure.core.credential.TokenRequestContext
+import com.azure.core.management.AzureEnvironment
+import com.azure.core.management.profile.AzureProfile
 import com.azure.identity.DefaultAzureCredentialBuilder
 import com.google.auth.oauth2.{AccessToken, OAuth2Credentials}
 import com.typesafe.config.Config
 import common.validation.ErrorOr.ErrorOr
 import net.ceedubs.ficus.Ficus._
 
+import java.time.{Duration => jDuration}
+import java.util.concurrent.TimeUnit
 import scala.concurrent.duration._
+import scala.util.{Failure, Success, Try}
 
 /**
   * This trait allows us to abstract away different token attainment strategies
@@ -45,12 +50,36 @@ object GoogleDrsCredentials {
 
 case class AzureDrsCredentials(identityClientId: Option[String]) extends DrsCredentials {
 
+  final val tokenAcquisitionTimeout = new jDuration(30, TimeUnit.SECONDS)
+
+  val azureProfile = new AzureProfile(AzureEnvironment.AZURE)
+  val tokenScope = "https://management.azure.com/.default"
+
+  def tokenRequestContext: TokenRequestContext = {
+    val trc = new TokenRequestContext()
+    trc.addScopes(tokenScope)
+    trc
+  }
+
+  def defaultCredentialBuilder: DefaultAzureCredentialBuilder =
+    new DefaultAzureCredentialBuilder()
+      .authorityHost(azureProfile.getEnvironment.getActiveDirectoryEndpoint)
+
   def getAccessToken: ErrorOr[String] = {
-    val credentials = identityClientId.foldLeft(new DefaultAzureCredentialBuilder()) {
+    val credentials = identityClientId.foldLeft(defaultCredentialBuilder) {
       (builder, clientId) => builder.managedIdentityClientId(clientId)
     }.build()
-    val tokenRequestContext = new TokenRequestContext()
-    tokenRequestContext.addScopes(".default")
-    credentials.getToken(tokenRequestContext).block().getToken.validNel
+
+    Try(
+      credentials
+        .getToken(tokenRequestContext)
+        .block(tokenAcquisitionTimeout)
+    ) match {
+      case Success(token) => token.getToken.validNel
+      case Failure(error) =>
+        Option(error.getMessage)  // it's possible the message is null
+          .getOrElse(error.toString)
+          .invalidNel
+    }
   }
 }