From 1192cf1750a7e156f4bf7d7f7d92b1561cac69f7 Mon Sep 17 00:00:00 2001
From: avishnus <avishnus@visa.com>
Date: Mon, 21 Oct 2024 23:20:22 +0530
Subject: [PATCH] Passing full username to session

---
 .../apache/kyuubi/service/TFrontendService.scala  | 14 ++++++++------
 .../server/KyuubiTHttpFrontendService.scala       | 15 ++++++++-------
 2 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/service/TFrontendService.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/service/TFrontendService.scala
index 801236f66b2..a9af8d10550 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/service/TFrontendService.scala
+++ b/kyuubi-common/src/main/scala/org/apache/kyuubi/service/TFrontendService.scala
@@ -142,16 +142,16 @@ abstract class TFrontendService(name: String)
    * The real user is the user used for session authentication.
    * The session user is the proxy user if proxy user is provided, otherwise is the real user.
    */
-  protected def getRealUserAndSessionUser(req: TOpenSessionReq): (String, String) = {
-    val realUser: String =
-      ServiceUtils.getShortName(authFactory.getRemoteUser.getOrElse(req.getUsername))
+  protected def getRealUserAndSessionUser(req: TOpenSessionReq): (String, String, String) = {
+    val fullUsername: String = authFactory.getRemoteUser.getOrElse(req.getUsername)
+    val realUser: String = ServiceUtils.getShortName(fullUsername)
     val sessionUser =
       if (req.getConfiguration == null) {
         realUser
       } else {
         getProxyUser(req.getConfiguration, authFactory.getIpAddress.orNull, realUser)
       }
-    realUser -> sessionUser
+    (fullUsername, realUser, sessionUser)
   }
 
   protected def getIpAddress: String = {
@@ -166,14 +166,16 @@ abstract class TFrontendService(name: String)
   protected def getSessionHandle(req: TOpenSessionReq, res: TOpenSessionResp): SessionHandle = {
     val protocol = getMinVersion(SERVER_VERSION, req.getClient_protocol)
     res.setServerProtocolVersion(protocol)
-    val (realUser, sessionUser) = getRealUserAndSessionUser(req)
+    val (fullUsername, realUser, sessionUser) = getRealUserAndSessionUser(req)
     val ipAddress = getIpAddress
     val configuration =
       Map(KYUUBI_CLIENT_IP_KEY -> ipAddress, KYUUBI_SERVER_IP_KEY -> serverAddr.getHostAddress) ++
         Option(req.getConfiguration).map(_.asScala.toMap).getOrElse(Map.empty[String, String]) ++
         Map(
           KYUUBI_SESSION_CONNECTION_URL_KEY -> connectionUrl,
-          KYUUBI_SESSION_REAL_USER_KEY -> realUser)
+          KYUUBI_SESSION_REAL_USER_KEY -> realUser,
+          "kyuubi.session.full.user" -> fullUsername
+        ) // Add full username here
     val sessionHandle = be.openSession(
       protocol,
       sessionUser,
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/KyuubiTHttpFrontendService.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/KyuubiTHttpFrontendService.scala
index b99f50310f9..2a871f3b58b 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/KyuubiTHttpFrontendService.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/KyuubiTHttpFrontendService.scala
@@ -335,18 +335,19 @@ final class KyuubiTHttpFrontendService(
     }
   }
 
-  override protected def getRealUserAndSessionUser(req: TOpenSessionReq): (String, String) = {
-    val realUser = getShortName(Option(AuthenticationFilter.getUserName)
-      .getOrElse(req.getUsername))
+  override protected def getRealUserAndSessionUser(req: TOpenSessionReq)
+      : (String, String, String) = {
+    val fullUsername = Option(SessionManager.getUserName).getOrElse(req.getUsername)
+    val realUser = getShortName(fullUsername)
     // using the remote ip address instead of that in proxy http header for authentication
-    val ipAddress: String = AuthenticationFilter.getUserIpAddress
-    val sessionUser: String = if (req.getConfiguration == null) {
+    val ipAddress: String = SessionManager.getIpAddress
+    val sessionUser = if (req.getConfiguration == null) {
       realUser
     } else {
       getProxyUser(req.getConfiguration, ipAddress, realUser)
     }
-    debug(s"Client's real user: $realUser, session user: $sessionUser")
-    realUser -> sessionUser
+    debug(s"Client's full user: $fullUsername, real user: $realUser, session user: $sessionUser")
+    (fullUsername, realUser, sessionUser)
   }
 
   private def getShortName(userName: String): String = {