fix: Refresh Fabric authentication during AI Functions retries

cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/CognitiveServiceBase.scala

@@ -329,6 +329,8 @@ object URLEncodingUtils {
 }
 
 private[ml] object ServiceAuthHeaders {
+  private[ml] case class Resolution(headers: Map[String, String], usesFabricFallback: Boolean)
+
   private[ml] def nonBlank(value: String): Boolean = value != null && value.trim.nonEmpty
 
   // Normalize a header map from any caller (a writer-supplied java.util.HashMap included) into a
@@ -360,15 +362,15 @@ private[ml] object ServiceAuthHeaders {
       .headOption
       .map { case (_, value) => canonicalName -> value }
 
-  def build(subscriptionKey: Option[String],
-            subscriptionKeyHeaderName: String,
-            aadHeaderName: String,
-            aadToken: Option[String],
-            customAuthHeader: Option[String],
-            customHeaders: Option[Map[String, String]],
-            fabricFallbackAuthHeader: => Option[String],
-            telemHeaders: Option[Map[String, String]],
-            contentType: Option[String]): Map[String, String] = {
+  def resolve(subscriptionKey: Option[String],
+              subscriptionKeyHeaderName: String,
+              aadHeaderName: String,
+              aadToken: Option[String],
+              customAuthHeader: Option[String],
+              customHeaders: Option[Map[String, String]],
+              fabricFallbackAuthHeader: => Option[String],
+              telemHeaders: Option[Map[String, String]],
+              contentType: Option[String]): Resolution = {
     val providedCustomHeaders = sanitizeHeaderMap(customHeaders)
 
     // Header names that carry credentials in this context, compared case-insensitively.
@@ -383,13 +385,18 @@ private[ml] object ServiceAuthHeaders {
     // higher-priority source above is absent. A fallback that acquires a token (and may throw) is
     // therefore never run while a subscription key, AAD token, explicit custom-auth header, or an
     // embedded customHeaders credential is present.
-    val authHeader: Option[(String, String)] = subscriptionKey.filter(nonBlank)
+    val explicitAuthHeader: Option[(String, String)] = subscriptionKey.filter(nonBlank)
       .map(value => subscriptionKeyHeaderName -> value)
       .orElse(aadToken.filter(nonBlank).map(value => aadHeaderName -> ("Bearer " + value)))
       .orElse(customAuthHeader.filter(nonBlank).map(value => aadHeaderName -> value))
       .orElse(embeddedCredential(providedCustomHeaders, subscriptionKeyHeaderName))
       .orElse(embeddedCredential(providedCustomHeaders, aadHeaderName))
-      .orElse(fabricFallbackAuthHeader.filter(nonBlank).map(value => aadHeaderName -> value))
+    val fabricAuthHeader = if (explicitAuthHeader.isDefined) {
+      None
+    } else {
+      fabricFallbackAuthHeader.filter(nonBlank).map(value => aadHeaderName -> value)
+    }
+    val authHeader = explicitAuthHeader.orElse(fabricAuthHeader)
 
     // Generic headers never carry auth: strip api-key/Authorization entries (any casing) so they
     // can neither override the resolved credential nor duplicate it under a different case.
@@ -412,7 +419,30 @@ private[ml] object ServiceAuthHeaders {
     }
     contentType.filterNot(StringUtils.isEmpty).foreach(value => headers += ("Content-Type" -> value))
 
-    new scala.collection.immutable.TreeMap[String, String]() ++ headers
+    Resolution(
+      new scala.collection.immutable.TreeMap[String, String]() ++ headers,
+      usesFabricFallback = fabricAuthHeader.isDefined)
+  }
+
+  def build(subscriptionKey: Option[String],
+            subscriptionKeyHeaderName: String,
+            aadHeaderName: String,
+            aadToken: Option[String],
+            customAuthHeader: Option[String],
+            customHeaders: Option[Map[String, String]],
+            fabricFallbackAuthHeader: => Option[String],
+            telemHeaders: Option[Map[String, String]],
+            contentType: Option[String]): Map[String, String] = {
+    resolve(
+      subscriptionKey,
+      subscriptionKeyHeaderName,
+      aadHeaderName,
+      aadToken,
+      customAuthHeader,
+      customHeaders,
+      fabricFallbackAuthHeader,
+      telemHeaders,
+      contentType).headers
   }
 }
 
@@ -503,12 +533,22 @@ trait HasCognitiveServiceInput extends HasURL with HasSubscriptionKey with HasAA
     getValueOpt(row, customHeaders)
   }
 
+  protected def supportsImplicitFabricAuthRetry: Boolean = false
+
   protected def addHeaders(req: HttpRequestBase,
                            row: Row,
                            addContentType: Boolean = true): Unit = {
 
-    val headers = getHeaders(row, addContentType)
-    headers.foreach { case (headerName, headerValue) => req.addHeader(headerName, headerValue) }
+    val resolution = resolveServiceAuthHeaders(row, addContentType, getFabricFallbackAuthHeader(row))
+    req.removeHeaders(HTTPRequestData.FabricAuthMarkerHeader)
+    resolution.headers.foreach { case (headerName, headerValue) =>
+      if (!HTTPRequestData.isFabricAuthMarker(headerName)) {
+        req.addHeader(headerName, headerValue)
+      }
+    }
+    if (resolution.usesFabricFallback && supportsImplicitFabricAuthRetry) {
+      req.addHeader(HTTPRequestData.FabricAuthMarkerHeader, "true")
+    }
   }
 
   // Returns a list of key-value pairs representing the headers
@@ -523,7 +563,14 @@ trait HasCognitiveServiceInput extends HasURL with HasSubscriptionKey with HasAA
   private[ml] def buildServiceAuthHeaders(row: Row,
                                           addContentType: Boolean,
                                           fabricFallbackAuthHeader: => Option[String]): Map[String, String] = {
-    ServiceAuthHeaders.build(
+    resolveServiceAuthHeaders(row, addContentType, fabricFallbackAuthHeader).headers
+  }
+
+  private def resolveServiceAuthHeaders(
+      row: Row,
+      addContentType: Boolean,
+      fabricFallbackAuthHeader: => Option[String]): ServiceAuthHeaders.Resolution = {
+    ServiceAuthHeaders.resolve(
       getValueOpt(row, subscriptionKey),
       subscriptionKeyHeaderName,
       aadHeaderName,

cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAI.scala

@@ -577,6 +577,16 @@ trait HasOpenAIFabricHeaders extends HasCognitiveServiceInput {
     runningOnFabric && usingDefaultOpenAIEndpoint && !hasCustomUrlRoot
   }
 
+  override protected def supportsImplicitFabricAuthRetry: Boolean = usingImplicitFabricEndpoint
+
+  abstract override protected def getFabricFallbackAuthHeader(row: Row): Option[String] = {
+    if (usingImplicitFabricEndpoint) {
+      super.getFabricFallbackAuthHeader(row)
+    } else {
+      None
+    }
+  }
+
   abstract override protected def getCustomHeaders(row: Row): Option[Map[String, String]] = {
     val headers = super.getCustomHeaders(row)
     if (usingImplicitFabricEndpoint) {

cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIFabricAuthRetrySuite.scala

@@ -0,0 +1,61 @@
+// Copyright (C) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License. See LICENSE in project root for information.
+
+package com.microsoft.azure.synapse.ml.services.openai
+
+import com.microsoft.azure.synapse.ml.core.test.base.TestBase
+import com.microsoft.azure.synapse.ml.io.http.HTTPRequestData
+import org.apache.http.client.methods.HttpPost
+import org.apache.spark.sql.Row
+
+class OpenAIFabricAuthRetrySuite extends TestBase {
+
+  private class AuthProvenanceProbe(fallbackAuthHeader: Option[String]) extends OpenAIChatCompletion {
+    var fallbackCalls = 0
+
+    override protected[openai] def runningOnFabric: Boolean = true
+
+    override protected[openai] def usingDefaultOpenAIEndpoint: Boolean = true
+
+    override protected def getFabricFallbackAuthHeader(row: Row): Option[String] = {
+      fallbackCalls += 1
+      fallbackAuthHeader
+    }
+
+    def requestData: HTTPRequestData = {
+      val request = new HttpPost("https://example.test/openai")
+      addHeaders(request, Row.empty, addContentType = false)
+      new HTTPRequestData(request)
+    }
+  }
+
+  test("only implicit Fabric authentication marks a request for refresh") {
+    val implicitProbe = new AuthProvenanceProbe(Some("MwcToken implicit"))
+    val implicitRequestData = implicitProbe.requestData
+    val implicitHttpRequest = implicitRequestData.toHTTPCore
+
+    assert(implicitProbe.fallbackCalls === 1)
+    assert(implicitRequestData.usesFabricAuth)
+    assert(implicitHttpRequest.getFirstHeader("Authorization").getValue === "MwcToken implicit")
+    assert(Option(implicitHttpRequest.getFirstHeader(HTTPRequestData.FabricAuthMarkerHeader)).isEmpty)
+
+    val explicitProbe = new AuthProvenanceProbe(Some("MwcToken implicit"))
+      .setCustomAuthHeader("MwcToken explicit")
+      .setCustomHeaders(Map(HTTPRequestData.FabricAuthMarkerHeader -> "true"))
+      .asInstanceOf[AuthProvenanceProbe]
+    val explicitRequestData = explicitProbe.requestData
+
+    assert(explicitProbe.fallbackCalls === 0)
+    assert(!explicitRequestData.usesFabricAuth)
+    assert(explicitRequestData.toHTTPCore.getFirstHeader("Authorization").getValue === "MwcToken explicit")
+
+    val apiKeyProbe = new AuthProvenanceProbe(Some("MwcToken implicit"))
+      .setSubscriptionKey("explicit-key")
+      .asInstanceOf[AuthProvenanceProbe]
+    val apiKeyRequestData = apiKeyProbe.requestData
+
+    assert(apiKeyProbe.fallbackCalls === 0)
+    assert(!apiKeyRequestData.usesFabricAuth)
+    assert(apiKeyRequestData.toHTTPCore.getFirstHeader("api-key").getValue === "explicit-key")
+  }
+}

core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/FabricClient.scala

@@ -6,9 +6,11 @@ package com.microsoft.azure.synapse.ml.fabric
 import spray.json.DefaultJsonProtocol.{StringJsonFormat, mapFormat}
 import spray.json._
 
-import java.net.{MalformedURLException, URL}
-import java.util.UUID
+import java.net.{MalformedURLException, URI, URL}
+import java.util.{Locale, UUID}
+import scala.collection.concurrent.TrieMap
 import scala.io.Source
+import scala.util.control.NonFatal
 
 object FabricClient extends RESTUtils {
   private val WorkloadEndpointTypeML = "ML";
@@ -19,6 +21,8 @@ object FabricClient extends RESTUtils {
   private val ContextFilePath = "/home/trusted-service-user/.trident-context";
   private val SparkConfPath = "/opt/spark/conf/spark-defaults.conf";
   private val ClusterInfoPath = "/opt/health-agent/conf/cluster-info.json";
+  private val CognitiveMwcRefreshLocks = TrieMap.empty[(String, String), AnyRef]
+  private val AmbiguousPathEncoding = "(?i)%(?:25)*(?:2e|2f|5c)".r
 
   lazy val CapacityID: Option[String] = getCapacityID;
   lazy val WorkspaceID: Option[String] = getWorkspaceID;
@@ -195,4 +199,69 @@ object FabricClient extends RESTUtils {
   def getCognitiveMWCTokenAuthHeader: String = {
     TokenLibrary.getCognitiveMwcTokenAuthHeader(WorkspaceID.getOrElse(""), ArtifactID.getOrElse(""))
   }
+
+  private[ml] def isEndpointUnder(requestUrl: String, endpointRoot: String): Boolean = {
+    try {
+      val request = new URL(requestUrl)
+      val root = new URL(endpointRoot)
+      val requestPort = if (request.getPort >= 0) request.getPort else request.getDefaultPort
+      val rootPort = if (root.getPort >= 0) root.getPort else root.getDefaultPort
+      val requestPath = trustedPath(new URI(requestUrl).getRawPath)
+      val rootPath = trustedPath(new URI(endpointRoot).getRawPath)
+      requestPath.exists(path =>
+        rootPath.exists(rootPrefix =>
+          request.getProtocol.equalsIgnoreCase("https") &&
+            root.getProtocol.equalsIgnoreCase("https") &&
+            request.getHost.equalsIgnoreCase(root.getHost) &&
+            requestPort == rootPort &&
+            path.startsWith(rootPrefix)))
+    } catch {
+      case _: MalformedURLException | _: java.net.URISyntaxException => false
+    }
+  }
+
+  private def trustedPath(rawPath: String): Option[String] = {
+    Option(rawPath).flatMap { path =>
+      val lowerPath = path.toLowerCase(Locale.ROOT)
+      val hasAmbiguousEncoding = AmbiguousPathEncoding.findFirstIn(lowerPath).nonEmpty
+      val normalizedPath = path.replaceAll("/+", "/")
+      val hasDotSegment = normalizedPath.split("/", -1).exists(segment => segment == "." || segment == "..")
+      if (hasAmbiguousEncoding || path.contains("\\") || hasDotSegment) None else Some(normalizedPath)
+    }
+  }
+
+  private[ml] def isOpenAIEndpoint(requestUrl: String): Boolean = {
+    try {
+      isEndpointUnder(requestUrl, MLWorkloadEndpointOpenAI)
+    } catch {
+      case NonFatal(_) => false
+    }
+  }
+
+  def refreshCognitiveMWCTokenAuthHeader(rejectedAuthHeader: String): String = {
+    val workspaceId = WorkspaceID.getOrElse("")
+    val artifactId = ArtifactID.getOrElse("")
+    val refreshLock = CognitiveMwcRefreshLocks.getOrElseUpdate((workspaceId, artifactId), new Object())
+    refreshAuthHeader(
+      rejectedAuthHeader,
+      refreshLock,
+      () => TokenLibrary.getCognitiveMwcTokenAuthHeader(workspaceId, artifactId),
+      () => TokenLibrary.invalidateSparkMwcToken(workspaceId, artifactId))
+  }
+
+  private[ml] def refreshAuthHeader(
+      rejectedAuthHeader: String,
+      refreshLock: AnyRef,
+      getCurrentAuthHeader: () => String,
+      invalidateAuthHeader: () => Unit): String = {
+    refreshLock.synchronized {
+      val currentAuthHeader = getCurrentAuthHeader()
+      if (currentAuthHeader != rejectedAuthHeader) {
+        currentAuthHeader
+      } else {
+        invalidateAuthHeader()
+        getCurrentAuthHeader()
+      }
+    }
+  }
 }