From e2ebca6ac5cad1328d03657be45fe545148ac787 Mon Sep 17 00:00:00 2001
From: Daniel Kleveros <daniel.kleveros@walmart.com>
Date: Sun, 24 Sep 2023 21:38:54 -0700
Subject: [PATCH] using sub from profile instead of token

---
 .../plugins/pfedsso/JwtAuthenticator.java     | 17 +-----------
 .../server/plugins/pfedsso/SsoClient.java     |  5 +++-
 .../server/plugins/pfedsso/SsoHandler.java    | 26 ++++++++++---------
 3 files changed, 19 insertions(+), 29 deletions(-)

diff --git a/server/plugins/pfed-sso/src/main/java/com/walmartlabs/concord/server/plugins/pfedsso/JwtAuthenticator.java b/server/plugins/pfed-sso/src/main/java/com/walmartlabs/concord/server/plugins/pfedsso/JwtAuthenticator.java
index 754400733f..7e6acc80ab 100644
--- a/server/plugins/pfed-sso/src/main/java/com/walmartlabs/concord/server/plugins/pfedsso/JwtAuthenticator.java
+++ b/server/plugins/pfed-sso/src/main/java/com/walmartlabs/concord/server/plugins/pfedsso/JwtAuthenticator.java
@@ -113,22 +113,7 @@ public boolean isTokenValid(String token, String nonce) {
         }
     }
 
-    /**
-     * Validates the token and returns the corresponding user login.
-     *
-     * @param token the JWT
-     * @return corresponding user login or <code>null</code> if the JWT is invalid
-     */
-    public String validateTokenAndGetLogin(String token) {
-        Map<String, Object> claims = validateTokenAndGetClaims(token);
-        if (claims == null) {
-            return null;
-        }
-        return (String) claims.get("sub");
-    }
-
-    private Map<String, Object> validateTokenAndGetClaims(String token) {
-
+    public Map<String, Object> validateTokenAndGetClaims(String token) {
         try {
             JWT jwt = validateToken(token);
             if (jwt == null) {
diff --git a/server/plugins/pfed-sso/src/main/java/com/walmartlabs/concord/server/plugins/pfedsso/SsoClient.java b/server/plugins/pfed-sso/src/main/java/com/walmartlabs/concord/server/plugins/pfedsso/SsoClient.java
index ed53348bc8..f07a303527 100644
--- a/server/plugins/pfed-sso/src/main/java/com/walmartlabs/concord/server/plugins/pfedsso/SsoClient.java
+++ b/server/plugins/pfed-sso/src/main/java/com/walmartlabs/concord/server/plugins/pfedsso/SsoClient.java
@@ -126,7 +126,7 @@ public String getTokenSigningKey() throws IOException {
         }
     }
 
-    public Profile getUserProfileFromRefreshToken(String refreshToken) throws IOException {
+    public Profile getUserProfileByRefreshToken(String refreshToken) throws IOException {
         Token token = getTokenByRefreshToken(refreshToken);
         return getProfile(token.accessToken());
     }
@@ -240,6 +240,9 @@ public interface Token {
     @JsonIgnoreProperties(ignoreUnknown = true)
     public interface Profile {
 
+        @JsonProperty("sub")
+        String sub();
+
         @JsonProperty("sAMAccountName")
         String userId();
 
diff --git a/server/plugins/pfed-sso/src/main/java/com/walmartlabs/concord/server/plugins/pfedsso/SsoHandler.java b/server/plugins/pfed-sso/src/main/java/com/walmartlabs/concord/server/plugins/pfedsso/SsoHandler.java
index 7a68af83d4..fe2b5a1a54 100644
--- a/server/plugins/pfed-sso/src/main/java/com/walmartlabs/concord/server/plugins/pfedsso/SsoHandler.java
+++ b/server/plugins/pfed-sso/src/main/java/com/walmartlabs/concord/server/plugins/pfedsso/SsoHandler.java
@@ -57,30 +57,32 @@ public AuthenticationToken createToken(ServletRequest request, ServletResponse r
         HttpServletRequest req = WebUtils.toHttp(request);
 
         String bearerToken = extractTokenFromRequest(req);
-        String incomingToken = bearerToken != null ? bearerToken : SsoCookies.getTokenCookie(req);
+        String token = bearerToken != null ? bearerToken : SsoCookies.getTokenCookie(req);
 
-        if (incomingToken == null) {
+        if (token == null) {
             return null;
         }
 
-        String login = jwtAuthenticator.validateTokenAndGetLogin(incomingToken);
-        if (login == null) {
+        if (!jwtAuthenticator.isTokenValid(token)) {
             return null;
         }
 
-        String[] as = parseDomain(login);
-
-        SsoClient.Profile profile;
         try {
-            profile = bearerToken != null ? ssoClient.getProfile(bearerToken) : ssoClient.getUserProfileFromRefreshToken(incomingToken);
+            SsoClient.Profile profile = bearerToken != null ? ssoClient.getProfile(bearerToken) :
+                    ssoClient.getUserProfileByRefreshToken(SsoCookies.getRefreshCookie(req));
+
+            if (profile == null) {
+                return null;
+            }
+
+            String[] as = parseDomain(profile.sub());
+
+            return new SsoToken(as[0], as[1], profile.displayName(), profile.mail(), profile.userPrincipalName(), profile.nameInNamespace(), profile.groups());
 
         } catch (IOException e) {
+
             return null;
         }
-        if (profile == null) {
-            return null;
-        }
-        return new SsoToken(as[0], as[1], profile.displayName(), profile.mail(), profile.userPrincipalName(), profile.nameInNamespace(), profile.groups());
     }
 
     @Override