aajisaka commented on code in PR #7483:
URL: https://github.com/apache/kyuubi/pull/7483#discussion_r3323124519


##########
kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/EngineUIProxyServlet.scala:
##########
@@ -52,17 +70,125 @@ private[api] class EngineUIProxyServlet extends 
ProxyServlet with Logging {
       clientRequest: HttpServletRequest,
       proxyRequest: Request): Unit = {
     val requestURI = clientRequest.getRequestURI
-    extractTargetAddress(requestURI).foreach { case (host, port) =>
+    extractTargetAddress(requestURI).filter(isRegistered).foreach { case 
(host, port) =>
       // SPARK-24209: Knox uses X-Forwarded-Context to notify the application 
the base path
       proxyRequest.header("X-Forwarded-Context", s"/engine-ui/$host:$port")
     }
     super.addXForwardedHeaders(clientRequest, proxyRequest)
   }
 
+  private def isRegistered(target: (String, Int)): Boolean = {
+    !filterEnabled || {
+      try {
+        currentRegisteredEngineUrls().contains(normalizeHostPort(target))
+      } catch {
+        case NonFatal(e) =>
+          warn(s"Failed to validate Engine UI target 
${target._1}:${target._2}", e)
+          false
+      }
+    }
+  }
+
+  private def currentRegisteredEngineUrls(): Set[(String, Int)] = {
+    val now = System.nanoTime()
+    val cacheTtlNanos = EngineUIProxyServlet.RegisteredEngineUrlsCacheTtlNanos
+    if (now - cachedAtNanos > cacheTtlNanos) {
+      synchronized {
+        val current = System.nanoTime()
+        if (current - cachedAtNanos > cacheTtlNanos) {
+          cachedRegisteredEngineUrls = registeredEngineUrls()
+          cachedAtNanos = current
+        }
+      }
+    }
+    cachedRegisteredEngineUrls
+  }
+
   private val r = "^/engine-ui/([^/:]+):(\\d+)/?.*".r
-  private def extractTargetAddress(requestURI: String): Option[(String, Int)] =
+  private[api] def extractTargetAddress(requestURI: String): Option[(String, 
Int)] =
     requestURI match {
       case r(host, port) => Some(host -> port.toInt)
       case _ => None
     }
+
+  private def normalizeHostPort(hostPort: (String, Int)): (String, Int) = {
+    hostPort._1.toLowerCase(Locale.ROOT) -> hostPort._2
+  }
+}
+
+private[api] object EngineUIProxyServlet extends Logging {
+
+  // Engine UI pages load many assets through the proxy. Avoid scanning 
discovery for every request.
+  private val RegisteredEngineUrlsCacheTtlNanos = TimeUnit.SECONDS.toNanos(10)
+
+  def apply(fe: KyuubiRestFrontendService): EngineUIProxyServlet = {
+    new EngineUIProxyServlet(
+      fe.getConf.get(KyuubiConf.FRONTEND_REST_ENGINE_UI_PROXY_FILTER_ENABLED),
+      () => registeredEngineUrls(fe))
+  }
+
+  private[api] def registeredEngineUrls(fe: KyuubiRestFrontendService): 
Set[(String, Int)] = {
+    DiscoveryClientProvider.withDiscoveryClient(fe.getConf) { discoveryClient 
=>
+      registeredEngineUrls(discoveryClient, fe.getConf.get(HA_NAMESPACE))
+    }
+  }
+
+  private[api] def registeredEngineUrls(

Review Comment:
   Probably we might need to set `*` by default for compatibility, but I feel 
it doesn't change the situation. Probably almost all the users rely on the 
default value.
   
   Instead, can we set the default to none? and if the host is not in the 
allowlist, just redirect it. The traffic directly goes to the engine URL. What 
do you think?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to