Fix username attributes for sso login

Aiven-Open · Aug 21, 2024 · 6f8d04c · 6f8d04c
1 parent 367af55
commit 6f8d04c
Show file tree

Hide file tree

Showing 7 changed files with 93 additions and 28 deletions.
diff --git a/core/src/main/java/io/aiven/klaw/auth/KwAuthenticationSuccessHandler.java b/core/src/main/java/io/aiven/klaw/auth/KwAuthenticationSuccessHandler.java
@@ -31,6 +31,9 @@ public class KwAuthenticationSuccessHandler extends SavedRequestAwareAuthenticat
   @Value("${klaw.ad.username.attribute:preferred_username}")
   private String preferredUsernameAttribute;
 
+  @Value("${klaw.ad.email.attribute:preferred_username}")
+  private String emailAttribute;
+
   @Autowired HandleDbRequestsJdbc handleDbRequests;
 
   @Override
@@ -53,7 +56,8 @@ public String getRedirectPage(HttpServletRequest request, Authentication authent
     if (quickStartEnabled
         && handleDbRequests
             .getUsersInfo(
-                UtilMethods.getUserName(authentication.getPrincipal(), preferredUsernameAttribute))
+                UtilMethods.getUserName(
+                    authentication.getPrincipal(), preferredUsernameAttribute, emailAttribute))
             .getRole()
             .equals(KwConstants.USER_ROLE)) {
       return coralTopicsUri;
@@ -63,7 +67,8 @@ public String getRedirectPage(HttpServletRequest request, Authentication authent
         && UtilControllerService.isCoralBuilt
         && !handleDbRequests
             .getUsersInfo(
-                UtilMethods.getUserName(authentication.getPrincipal(), preferredUsernameAttribute))
+                UtilMethods.getUserName(
+                    authentication.getPrincipal(), preferredUsernameAttribute, emailAttribute))
             .getRole()
             .equals(KwConstants.SUPERADMIN_ROLE)) {
       return coralTopicsUri;

diff --git a/core/src/main/java/io/aiven/klaw/controller/ResourceClientController.java b/core/src/main/java/io/aiven/klaw/controller/ResourceClientController.java
@@ -33,6 +33,9 @@ public class ResourceClientController {
   @Value("${klaw.ad.username.attribute:preferred_username}")
   private String preferredUsernameAttribute;
 
+  @Value("${klaw.ad.email.attribute:preferred_username}")
+  private String emailAttribute;
+
   private static final String authorizationRequestBaseUri = "oauth2/authorize-client";
   Map<String, String> oauth2AuthenticationUrls = new HashMap<>();
   @Autowired private OAuth2AuthorizedClientService authorizedClientService;
@@ -64,10 +67,13 @@ private String checkAnonymousLogin(
     try {
       Object principal = SecurityContextHolder.getContext().getAuthentication().getPrincipal();
       DefaultOAuth2User defaultOAuth2User = (DefaultOAuth2User) principal;
+      String userName = (String) defaultOAuth2User.getAttributes().get(preferredUsernameAttribute);
+      if (userName == null) {
+        userName = (String) defaultOAuth2User.getAttributes().get(emailAttribute);
+      }
       OAuth2AuthorizedClient client =
           authorizedClientService.loadAuthorizedClient(
-              authentication.getAuthorizedClientRegistrationId(),
-              (String) defaultOAuth2User.getAttributes().get(preferredUsernameAttribute));
+              authentication.getAuthorizedClientRegistrationId(), userName);
       if (client == null) {
         return ("redirect:oauthLogin");
       }

diff --git a/core/src/main/java/io/aiven/klaw/helpers/UtilMethods.java b/core/src/main/java/io/aiven/klaw/helpers/UtilMethods.java
@@ -22,18 +22,21 @@
 
 @Slf4j
 public class UtilMethods {
-  public static String getUserName(Object principal, String preferredUsername) {
+  public static String getUserName(
+      Object principal, String preferredUserNameAttribute, String emailAttribute) {
     if (principal instanceof DefaultOAuth2User defaultOAuth2User) {
-      return (String) defaultOAuth2User.getAttributes().get(preferredUsername);
+      return Optional.ofNullable(
+              (String) defaultOAuth2User.getAttributes().get(preferredUserNameAttribute))
+          .orElse((String) defaultOAuth2User.getAttributes().get(emailAttribute));
     } else if (principal instanceof String) {
       return (String) principal;
     } else {
       return ((UserDetails) principal).getUsername();
     }
   }
 
-  public static String getUserName(String preferredUsername) {
-    return getUserName(getPrincipal(), preferredUsername);
+  public static String getUserName(String preferredUserNameAttribute, String emailAttribute) {
+    return getUserName(getPrincipal(), preferredUserNameAttribute, emailAttribute);
   }
 
   public static Object getPrincipal() {

diff --git a/core/src/main/java/io/aiven/klaw/service/CommonUtilsService.java b/core/src/main/java/io/aiven/klaw/service/CommonUtilsService.java
@@ -109,11 +109,16 @@ public Authentication getAuthentication() {
     return SecurityContextHolder.getContext().getAuthentication();
   }
 
+  public ArrayList<String> getUserNameAttributes() {
+    return new ArrayList<>(List.of(preferredUsernameAttribute, emailAttribute));
+  }
+
   String getAuthority(Object principal) {
     if (enableUserAuthorizationFromAD) {
       if (principal instanceof DefaultOAuth2User) {
         DefaultOAuth2User defaultOAuth2User = (DefaultOAuth2User) principal;
-        String userName = extractUserNameFromOAuthUser(defaultOAuth2User);
+        String userName = getUserName(defaultOAuth2User);
+
         return manageDatabase.getHandleDbRequests().getUsersInfo(userName).getRole();
       } else if (principal instanceof String) {
         return manageDatabase.getHandleDbRequests().getUsersInfo((String) principal).getRole();
@@ -138,25 +143,12 @@ String getAuthority(Object principal) {
     }
   }
 
-  public String extractUserNameFromOAuthUser(DefaultOAuth2User defaultOAuth2User) {
-    String preferredUsername =
-        (String) defaultOAuth2User.getAttributes().get(preferredUsernameAttribute);
-    String email = (String) defaultOAuth2User.getAttributes().get(emailAttribute);
-    String userName = null;
-    if (preferredUsername != null) {
-      userName = preferredUsername;
-    } else if (email != null) {
-      userName = email;
-    }
-    return userName;
-  }
-
   public String getUserName(Object principal) {
-    return UtilMethods.getUserName(principal, preferredUsernameAttribute);
+    return UtilMethods.getUserName(principal, preferredUsernameAttribute, emailAttribute);
   }
 
   public String getCurrentUserName() {
-    return UtilMethods.getUserName(preferredUsernameAttribute);
+    return UtilMethods.getUserName(preferredUsernameAttribute, emailAttribute);
   }
 
   public boolean isNotAuthorizedUser(Object principal, PermissionType permissionType) {

diff --git a/core/src/main/java/io/aiven/klaw/service/MailUtils.java b/core/src/main/java/io/aiven/klaw/service/MailUtils.java
@@ -42,7 +42,10 @@ public class MailUtils {
   private String kwAdminMailId;
 
   @Value("${klaw.ad.username.attribute:preferred_username}")
-  private String preferredUsername;
+  private String preferredUsernameAttribute;
+
+  @Value("${klaw.ad.email.attribute:preferred_username}")
+  private String emailAttribute;
 
   private static final String TOPIC_REQ_KEY = "klaw.mail.topicrequest.content";
   private static final String TOPIC_PROMOTION_REQ_KEY = "klaw.mail.topicpromotionrequest.content";
@@ -68,11 +71,11 @@ public class MailUtils {
   @Autowired private EmailService emailService;
 
   public String getUserName(Object principal) {
-    return UtilMethods.getUserName(principal, preferredUsername);
+    return UtilMethods.getUserName(principal, preferredUsernameAttribute, emailAttribute);
   }
 
   public String getCurrentUserName() {
-    return UtilMethods.getUserName(preferredUsername);
+    return UtilMethods.getUserName(preferredUsernameAttribute, emailAttribute);
   }
 
   void sendMail(

diff --git a/core/src/main/java/io/aiven/klaw/service/UiControllerLoginService.java b/core/src/main/java/io/aiven/klaw/service/UiControllerLoginService.java
@@ -225,7 +225,7 @@ public String checkAuth(
     if (abstractAuthenticationToken instanceof OAuth2AuthenticationToken) {
       DefaultOAuth2User defaultOAuth2User =
           (DefaultOAuth2User) SecurityContextHolder.getContext().getAuthentication().getPrincipal();
-      userName = commonUtilsService.extractUserNameFromOAuthUser(defaultOAuth2User);
+      userName = commonUtilsService.getUserName(defaultOAuth2User);
     } else if (abstractAuthenticationToken instanceof UsernamePasswordAuthenticationToken) {
       userName =
           ((UserDetails) SecurityContextHolder.getContext().getAuthentication().getPrincipal())

diff --git a/core/src/test/java/io/aiven/klaw/helpers/UtilMethodsTest.java b/core/src/test/java/io/aiven/klaw/helpers/UtilMethodsTest.java
@@ -0,0 +1,56 @@
+package io.aiven.klaw.helpers;
+
+import static org.junit.jupiter.api.Assertions.*;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import java.util.HashMap;
+import java.util.Map;
+import org.junit.jupiter.api.Test;
+import org.springframework.security.core.userdetails.UserDetails;
+import org.springframework.security.oauth2.core.user.DefaultOAuth2User;
+
+class UtilMethodsTest {
+
+  @Test
+  public void testGetUserNameFromOAuth2User_PreferredUserName() {
+    Map<String, Object> attributes = new HashMap<>();
+    attributes.put("preferred_username", "testUser");
+    attributes.put("email", "[email protected]");
+
+    DefaultOAuth2User oAuth2User = mock(DefaultOAuth2User.class);
+    when(oAuth2User.getAttributes()).thenReturn(attributes);
+
+    String result = UtilMethods.getUserName(oAuth2User, "preferred_username", "email");
+    assertEquals("testUser", result);
+  }
+
+  @Test
+  public void testGetUserNameFromOAuth2User_Email() {
+    Map<String, Object> attributes = new HashMap<>();
+    attributes.put("email", "[email protected]");
+
+    DefaultOAuth2User oAuth2User = mock(DefaultOAuth2User.class);
+    when(oAuth2User.getAttributes()).thenReturn(attributes);
+
+    String result = UtilMethods.getUserName(oAuth2User, "preferred_username", "email");
+    assertEquals("[email protected]", result);
+  }
+
+  @Test
+  public void testGetUserNameFromStringPrincipal() {
+    String principal = "testUser";
+
+    String result = UtilMethods.getUserName(principal, "preferred_username", "email");
+    assertEquals("testUser", result);
+  }
+
+  @Test
+  public void testGetUserNameFromUserDetails() {
+    UserDetails userDetails = mock(UserDetails.class);
+    when(userDetails.getUsername()).thenReturn("testUser");
+
+    String result = UtilMethods.getUserName(userDetails, "preferred_username", "email");
+    assertEquals("testUser", result);
+  }
+}