[SPARK-58893][UI] Fix Spark UI multi-cluster reverse proxy: propagate…

core/src/main/scala/org/apache/spark/SparkContext.scala

@@ -636,7 +636,9 @@ class SparkContext(config: SparkConf) extends Logging {
     }
 
     if (_conf.get(UI_REVERSE_PROXY)) {
-      val proxyUrl = _conf.get(UI_REVERSE_PROXY_URL).getOrElse("").stripSuffix("/")
+      val proxyUrl = _conf.get(UI_REVERSE_PROXY_URL)
+        .getOrElse(sys.props.getOrElse("spark.ui.proxyBase", ""))
+        .stripSuffix("/")
       System.setProperty("spark.ui.proxyBase", proxyUrl + "/proxy/" + _applicationId)
     }
     _ui.foreach(_.setAppId(_applicationId))

core/src/main/scala/org/apache/spark/deploy/master/Master.scala

@@ -1345,11 +1345,24 @@ private[deploy] class Master(
     desc.copy(command = desc.command.copy(arguments = arguments, javaOpts = javaOpts))
   }
 
+  private def maybeAddReverseProxyConfig(desc: DriverDescription): DriverDescription = {
+    if (!reverseProxy) return desc
+    conf.get(UI_REVERSE_PROXY_URL).map(_.stripSuffix("/")).filter(_.nonEmpty) match {
+      case Some(url) =>
+        val opt = s"-Dspark.ui.reverseProxyUrl=$url"
+        val javaOpts = desc.command.javaOpts
+          .filter(!_.startsWith("-Dspark.ui.reverseProxyUrl=")) :+ opt
+        desc.copy(command = desc.command.copy(javaOpts = javaOpts))
+      case None => desc
+    }
+  }
+
   private def createDriver(desc: DriverDescription): DriverInfo = {
     val now = System.currentTimeMillis()
     val date = new Date(now)
     val id = newDriverId(date)
-    new DriverInfo(now, id, maybeUpdateAppName(desc, id), date)
+    val updatedDesc = maybeAddReverseProxyConfig(maybeUpdateAppName(desc, id))
+    new DriverInfo(now, id, updatedDesc, date)
   }
 
   private def launchDriver(worker: WorkerInfo, driver: DriverInfo): Unit = {

core/src/main/scala/org/apache/spark/deploy/master/ui/MasterWebUI.scala

@@ -29,6 +29,7 @@ import org.apache.spark.internal.LogKeys.{HOSTS, NUM_REMOVED_WORKERS}
 import org.apache.spark.internal.config.DECOMMISSION_ENABLED
 import org.apache.spark.internal.config.UI.MASTER_UI_DECOMMISSION_ALLOW_MODE
 import org.apache.spark.internal.config.UI.UI_KILL_ENABLED
+import org.apache.spark.internal.config.UI.UI_REVERSE_PROXY_URL
 import org.apache.spark.ui.{SparkUI, WebUI}
 import org.apache.spark.ui.JettyUtils._
 import org.apache.spark.util.ArrayImplicits._
@@ -64,10 +65,15 @@ class MasterWebUI(
     addStaticHandler(MasterWebUI.STATIC_RESOURCE_DIR)
     addRenderLogHandler(this, master.conf)
     if (killEnabled) {
+      val killRedirectTarget = master.conf.get(UI_REVERSE_PROXY_URL)
+        .map(_.stripSuffix("/") + "/")
+        .getOrElse("/")
       attachHandler(createRedirectHandler(
-        "/app/kill", "/", masterPage.handleAppKillRequest, httpMethods = Set("POST")))
+        "/app/kill", killRedirectTarget, masterPage.handleAppKillRequest,
+        httpMethods = Set("POST")))
       attachHandler(createRedirectHandler(
-        "/driver/kill", "/", masterPage.handleDriverKillRequest, httpMethods = Set("POST")))
+        "/driver/kill", killRedirectTarget, masterPage.handleDriverKillRequest,
+        httpMethods = Set("POST")))
     }
     if (decommissionEnabled) {
       attachHandler(createServletHandler("/workers/kill", new HttpServlet {
@@ -97,7 +103,8 @@ class MasterWebUI(
   }
 
   def addProxy(): Unit = {
-    val handler = createProxyHandler(idToUiAddress)
+    val reverseProxyUrl = master.conf.get(UI_REVERSE_PROXY_URL).getOrElse("")
+    val handler = createProxyHandler(idToUiAddress, reverseProxyUrl)
     attachHandler(handler)
   }
 

core/src/main/scala/org/apache/spark/ui/JettyUtils.scala

@@ -182,7 +182,9 @@ private[spark] object JettyUtils extends Logging {
   }
 
   /** Create a handler for proxying request to Workers and Application Drivers */
-  def createProxyHandler(idToUiAddress: String => Option[String]): ServletContextHandler = {
+  def createProxyHandler(
+      idToUiAddress: String => Option[String],
+      reverseProxyUrl: String = ""): ServletContextHandler = {
     val servlet = new ProxyServlet {
       override def rewriteTarget(request: HttpServletRequest): String = {
         val path = request.getPathInfo
@@ -206,6 +208,25 @@ private[spark] object JettyUtils extends Logging {
           .orNull
       }
 
+      override def addProxyHeaders(
+          clientRequest: HttpServletRequest,
+          proxyRequest: org.eclipse.jetty.client.api.Request): Unit = {
+        super.addProxyHeaders(clientRequest, proxyRequest)
+        val path = clientRequest.getPathInfo
+        if (path != null) {
+          val prefixTrailingSlashIndex = path.indexOf('/', 1)
+          val prefix = if (prefixTrailingSlashIndex == -1) {
+            path
+          } else {
+            path.substring(0, prefixTrailingSlashIndex)
+          }
+          val existingContext = Option(clientRequest.getHeader("X-Forwarded-Context")).getOrElse("")
+          val contextPath = Option(clientRequest.getContextPath).getOrElse("")
+          val proxyContext = existingContext + reverseProxyUrl + contextPath + prefix
+          proxyRequest.headers(headers => headers.put("X-Forwarded-Context", proxyContext))
+        }
+      }
+
       override def newHttpClient(): HttpClient = {
         // SPARK-21176: Use the Jetty logic to calculate the number of selector threads (#CPUs/2),
         // but limit it to 8 max.
@@ -222,6 +243,14 @@ private[spark] object JettyUtils extends Logging {
           val newHeader = createProxyLocationHeader(headerValue, clientRequest,
             serverResponse.getRequest().getURI())
           if (newHeader != null) {
+            if (reverseProxyUrl.nonEmpty) {
+              val scheme = clientRequest.getScheme
+              val host = Option(clientRequest.getHeader("host")).getOrElse("")
+              val rootProxyPrefix = s"$scheme://$host/proxy/"
+              if (newHeader.startsWith(rootProxyPrefix)) {
+                return reverseProxyUrl + "/proxy/" + newHeader.substring(rootProxyPrefix.length)
+              }
+            }
             return newHeader
           }
         }

core/src/main/scala/org/apache/spark/ui/UIUtils.scala

@@ -189,7 +189,7 @@ private[spark] object UIUtils extends Logging {
   // Yarn has to go through a proxy so the base uri is provided and has to be on all links
   def uiRoot(request: HttpServletRequest): String = {
     // Knox uses X-Forwarded-Context to notify the application the base path
-    val knoxBasePath = Option(request.getHeader("X-Forwarded-Context"))
+    val knoxBasePath = Option(request).flatMap(r => Option(r.getHeader("X-Forwarded-Context")))
     // SPARK-11484 - Use the proxyBase set by the AM, if not found then use env.
     sys.props.get("spark.ui.proxyBase")
       .orElse(sys.env.get("APPLICATION_WEB_PROXY_BASE"))
@@ -201,7 +201,14 @@ private[spark] object UIUtils extends Logging {
       request: HttpServletRequest,
       basePath: String = "",
       resource: String = ""): String = {
-    uiRoot(request) + basePath + resource
+    val root = uiRoot(request).stripSuffix("/")
+    val cleanBase = if (basePath.startsWith("/")) basePath
+                    else if (basePath.nonEmpty) "/" + basePath
+                    else ""
+    val cleanResource = if (resource.startsWith("/")) resource
+                        else if (resource.nonEmpty) "/" + resource
+                        else ""
+    s"$root$cleanBase$cleanResource"
   }
 
   def commonHeaderNodes(request: HttpServletRequest): Seq[Node] = {

core/src/test/scala/org/apache/spark/SparkContextSuite.scala

@@ -1371,6 +1371,31 @@ class SparkContextSuite extends SparkFunSuite with LocalSparkContext with Eventu
     assert(msg.contains("Cannot use the keyword 'proxy' or 'history' in reverse proxy URL"))
   }
 
+  test("SPARK-58893: propagate UI_REVERSE_PROXY_URL to spark.ui.proxyBase in SparkContext") {
+    val conf = new SparkConf().setAppName("testReverseProxyBase")
+      .setMaster("local")
+      .set(UI_REVERSE_PROXY, true)
+      .set(UI_REVERSE_PROXY_URL, "http://proxyhost:8080/myprefix")
+    sc = new SparkContext(conf)
+    assert(System.getProperty("spark.ui.proxyBase") ===
+      s"http://proxyhost:8080/myprefix/proxy/${sc.applicationId}")
+  }
+
+  test("SPARK-58893: fallback spark.ui.proxyBase when UI_REVERSE_PROXY_URL is empty") {
+    val sysProxyBase = "http://proxyhost:8080/sysprefix"
+    System.setProperty("spark.ui.proxyBase", sysProxyBase)
+    try {
+      val conf = new SparkConf().setAppName("testReverseProxyBaseSys")
+        .setMaster("local")
+        .set(UI_REVERSE_PROXY, true)
+      sc = new SparkContext(conf)
+      assert(System.getProperty("spark.ui.proxyBase") ===
+        s"$sysProxyBase/proxy/${sc.applicationId}")
+    } finally {
+      System.clearProperty("spark.ui.proxyBase")
+    }
+  }
+
   test("SPARK-39957: ExitCode HEARTBEAT_FAILURE should be counted as network failure") {
     // This test is used to prove that driver will receive executorExitCode before onDisconnected
     // removes the executor. If the executor is removed by onDisconnected, the executor loss will be

core/src/test/scala/org/apache/spark/deploy/master/MasterSuite.scala

@@ -28,6 +28,7 @@ import org.apache.spark.deploy._
 import org.apache.spark.deploy.DeployMessages._
 import org.apache.spark.internal.config._
 import org.apache.spark.internal.config.Deploy._
+import org.apache.spark.internal.config.UI._
 import org.apache.spark.resource.ResourceProfile
 import org.apache.spark.rpc.{RpcAddress, RpcEndpoint, RpcEnv}
 
@@ -259,4 +260,16 @@ class MasterSuite extends MasterSuiteBase {
       makeMaster(conf)
     }
   }
+
+  test("SPARK-58893: Master injects reverseProxyUrl into driver javaOpts") {
+    val conf = new SparkConf()
+      .set(UI_REVERSE_PROXY, true)
+      .set(UI_REVERSE_PROXY_URL, "http://proxyhost:8080/path")
+    val master = makeMaster(conf)
+    val command = Command("mainClass", Seq.empty, Map.empty, Seq.empty, Seq.empty, Seq.empty)
+    val desc = DriverDescription("", 1, 1, false, command)
+    val result = master.invokePrivate(_maybeAddReverseProxyConfig(desc))
+    val opt = "-Dspark.ui.reverseProxyUrl=http://proxyhost:8080/path"
+    assert(result.command.javaOpts.contains(opt))
+  }
 }

core/src/test/scala/org/apache/spark/deploy/master/MasterSuiteBase.scala

@@ -444,6 +444,8 @@ trait MasterSuiteBase extends SparkFunSuite
   protected val _newApplicationId = PrivateMethod[String](Symbol("newApplicationId"))
   protected val _maybeUpdateAppName =
     PrivateMethod[DriverDescription](Symbol("maybeUpdateAppName"))
+  protected val _maybeAddReverseProxyConfig =
+    PrivateMethod[DriverDescription](Symbol("maybeAddReverseProxyConfig"))
   protected val _createApplication = PrivateMethod[ApplicationInfo](Symbol("createApplication"))
 
   protected val workerInfo = makeWorkerInfo(512, 10)

core/src/test/scala/org/apache/spark/deploy/master/ui/MasterWebUISuite.scala

@@ -32,7 +32,7 @@ import org.apache.spark.deploy.DeployMessages.{DecommissionWorkersOnHosts, KillD
 import org.apache.spark.deploy.DeployTestUtils._
 import org.apache.spark.deploy.master._
 import org.apache.spark.internal.config.DECOMMISSION_ENABLED
-import org.apache.spark.internal.config.UI.MASTER_UI_DECOMMISSION_ALLOW_MODE
+import org.apache.spark.internal.config.UI.{MASTER_UI_DECOMMISSION_ALLOW_MODE, UI_REVERSE_PROXY, UI_REVERSE_PROXY_URL}
 import org.apache.spark.rpc.{RpcEndpointRef, RpcEnv}
 import org.apache.spark.util.Utils
 
@@ -128,6 +128,42 @@ class MasterWebUISuite extends SparkFunSuite {
       denyWebUI.stop()
     }
   }
+
+  test("SPARK-58893: kill application redirect location with reverse proxy") {
+    val reverseProxyConf = new SparkConf()
+      .set(DECOMMISSION_ENABLED, true)
+      .set(UI_REVERSE_PROXY, true)
+      .set(UI_REVERSE_PROXY_URL, "http://proxyhost:8080/myproxy")
+    val mockMaster = mock(classOf[Master])
+    when(mockMaster.securityMgr).thenReturn(securityMgr)
+    when(mockMaster.conf).thenReturn(reverseProxyConf)
+    when(mockMaster.rpcEnv).thenReturn(rpcEnv)
+    when(mockMaster.self).thenReturn(masterEndpointRef)
+
+    val activeApp = new ApplicationInfo(
+      new Date().getTime, "app-proxy-0", createAppDesc(), new Date(), null, Int.MaxValue)
+    val appMap = HashMap[String, ApplicationInfo]((activeApp.id, activeApp))
+    when(mockMaster.idToApp).thenReturn(appMap)
+
+    val webUI = new MasterWebUI(mockMaster, 0)
+    try {
+      webUI.bind()
+      val url = s"http://${Utils.localHostNameForURI()}:${webUI.boundPort}/app/kill/"
+      val body = convPostDataToString(Map(("id", activeApp.id), ("terminate", "true")))
+      val conn = new URI(url).toURL.openConnection().asInstanceOf[HttpURLConnection]
+      conn.setInstanceFollowRedirects(false)
+      conn.setRequestMethod("POST")
+      conn.setDoOutput(true)
+      conn.setRequestProperty("Content-Type", "application/x-www-form-urlencoded")
+      val out = new DataOutputStream(conn.getOutputStream)
+      out.write(body.getBytes(StandardCharsets.UTF_8))
+      out.close()
+      assert(conn.getResponseCode === 302)
+      assert(conn.getHeaderField("Location") === "http://proxyhost:8080/myproxy/")
+    } finally {
+      webUI.stop()
+    }
+  }
 }
 
 object MasterWebUISuite {

core/src/test/scala/org/apache/spark/ui/UIUtilsSuite.scala

@@ -225,4 +225,20 @@ class UIUtilsSuite extends SparkFunSuite {
     assert(cell4 === <td>{"java.lang.RuntimeException"}{UIUtils.detailsUINode(isMultiline = true, e4)}</td>)
   }
   // scalastyle:on line.size.limit
+
+  test("SPARK-58893: prependBaseUri slash handling") {
+    try {
+      System.setProperty("spark.ui.proxyBase", "http://localhost:8080/foo/")
+      assert(UIUtils.prependBaseUri(null, "bar", "baz") ===
+        "http://localhost:8080/foo/bar/baz")
+      assert(UIUtils.prependBaseUri(null, "/bar", "/baz") ===
+        "http://localhost:8080/foo/bar/baz")
+      assert(UIUtils.prependBaseUri(null, "", "baz") ===
+        "http://localhost:8080/foo/baz")
+      assert(UIUtils.prependBaseUri(null, "", "") ===
+        "http://localhost:8080/foo")
+    } finally {
+      System.clearProperty("spark.ui.proxyBase")
+    }
+  }
 }