From 137dcf11e9940ef2f4a9dd075f0ac2df548e0d4c Mon Sep 17 00:00:00 2001
From: Cheng Pan <chengpan@apache.org>
Date: Thu, 14 Sep 2023 22:06:30 +0800
Subject: [PATCH] [KYUUBI #5289] RESTful API should always print audit log

### _Why are the changes needed?_

Previously, the RESTful audit log did not contain HTTP requests that threw non-AuthenticationException during the process.

### _How was this patch tested?_
- [ ] Add some test cases that check the changes thoroughly including negative and positive cases if possible

- [ ] Add screenshots for manual tests if appropriate

- [x] [Run test](https://kyuubi.readthedocs.io/en/master/contributing/code/testing.html#running-tests) locally before make a pull request

### _Was this patch authored or co-authored using generative AI tooling?_

No.

Closes #5289 from pan3793/auth-log.

Closes #5289

9e292793e [Cheng Pan] RESTful API should always print audit log

Authored-by: Cheng Pan <chengpan@apache.org>
Signed-off-by: Cheng Pan <chengpan@apache.org>
---
 .../authentication/AuthenticationFilter.scala | 44 +++++++++----------
 1 file changed, 22 insertions(+), 22 deletions(-)

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/http/authentication/AuthenticationFilter.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/http/authentication/AuthenticationFilter.scala
index 11e856a29..523d24907 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/http/authentication/AuthenticationFilter.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/http/authentication/AuthenticationFilter.scala
@@ -22,7 +22,7 @@ import javax.security.sasl.AuthenticationException
 import javax.servlet.{Filter, FilterChain, FilterConfig, ServletException, ServletRequest, ServletResponse}
 import javax.servlet.http.{HttpServletRequest, HttpServletResponse}
 
-import scala.collection.mutable.HashMap
+import scala.collection.mutable
 
 import org.apache.kyuubi.Logging
 import org.apache.kyuubi.config.KyuubiConf
@@ -35,7 +35,8 @@ class AuthenticationFilter(conf: KyuubiConf) extends Filter with Logging {
   import AuthenticationHandler._
   import AuthSchemes._
 
-  private[authentication] val authSchemeHandlers = new HashMap[AuthScheme, AuthenticationHandler]()
+  private[authentication] val authSchemeHandlers =
+    new mutable.HashMap[AuthScheme, AuthenticationHandler]()
 
   private[authentication] def addAuthHandler(authHandler: AuthenticationHandler): Unit = {
     authHandler.init(conf)
@@ -109,32 +110,31 @@ class AuthenticationFilter(conf: KyuubiConf) extends Filter with Logging {
     HTTP_PROXY_HEADER_CLIENT_IP_ADDRESS.set(
       httpRequest.getHeader(conf.get(FRONTEND_PROXY_HTTP_CLIENT_IP_HEADER)))
 
-    if (matchedHandler == null) {
-      debug(s"No auth scheme matched for url: ${httpRequest.getRequestURL}")
-      httpResponse.setStatus(HttpServletResponse.SC_UNAUTHORIZED)
-      AuthenticationAuditLogger.audit(httpRequest, httpResponse)
-      httpResponse.sendError(
-        HttpServletResponse.SC_UNAUTHORIZED,
-        s"No auth scheme matched for $authorization")
-    } else {
-      HTTP_AUTH_TYPE.set(matchedHandler.authScheme.toString)
-      try {
+    try {
+      if (matchedHandler == null) {
+        debug(s"No auth scheme matched for url: ${httpRequest.getRequestURL}")
+        httpResponse.setStatus(HttpServletResponse.SC_UNAUTHORIZED)
+        httpResponse.sendError(
+          HttpServletResponse.SC_UNAUTHORIZED,
+          s"No auth scheme matched for $authorization")
+      } else {
+        HTTP_AUTH_TYPE.set(matchedHandler.authScheme.toString)
         val authUser = matchedHandler.authenticate(httpRequest, httpResponse)
         if (authUser != null) {
           HTTP_CLIENT_USER_NAME.set(authUser)
           doFilter(filterChain, httpRequest, httpResponse)
         }
-        AuthenticationAuditLogger.audit(httpRequest, httpResponse)
-      } catch {
-        case e: AuthenticationException =>
-          httpResponse.setStatus(HttpServletResponse.SC_FORBIDDEN)
-          AuthenticationAuditLogger.audit(httpRequest, httpResponse)
-          HTTP_CLIENT_USER_NAME.remove()
-          HTTP_CLIENT_IP_ADDRESS.remove()
-          HTTP_PROXY_HEADER_CLIENT_IP_ADDRESS.remove()
-          HTTP_AUTH_TYPE.remove()
-          httpResponse.sendError(HttpServletResponse.SC_FORBIDDEN, e.getMessage)
       }
+    } catch {
+      case e: AuthenticationException =>
+        httpResponse.setStatus(HttpServletResponse.SC_FORBIDDEN)
+        HTTP_CLIENT_USER_NAME.remove()
+        HTTP_CLIENT_IP_ADDRESS.remove()
+        HTTP_PROXY_HEADER_CLIENT_IP_ADDRESS.remove()
+        HTTP_AUTH_TYPE.remove()
+        httpResponse.sendError(HttpServletResponse.SC_FORBIDDEN, e.getMessage)
+    } finally {
+      AuthenticationAuditLogger.audit(httpRequest, httpResponse)
     }
   }