extends the functionality of oauth in Stirling PDF

2024-05-12 19:58:34 +02:00
parent f2b7aeeb1c
commit 811c19e00d
49 changed files with 724 additions and 210 deletions
--- a/src/main/java/stirling/software/SPDF/config/security/CustomAuthenticationFailureHandler.java
+++ b/src/main/java/stirling/software/SPDF/config/security/CustomAuthenticationFailureHandler.java
@@ -3,10 +3,12 @@ package stirling.software.SPDF.config.security;
 import java.io.IOException;
 import java.util.Optional;

-import org.springframework.beans.factory.annotation.Autowired;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.springframework.security.authentication.BadCredentialsException;
 import org.springframework.security.authentication.LockedException;
 import org.springframework.security.core.AuthenticationException;
+import org.springframework.security.core.userdetails.UsernameNotFoundException;
 import org.springframework.security.web.authentication.SimpleUrlAuthenticationFailureHandler;
 import org.springframework.stereotype.Component;

@@ -18,9 +20,12 @@ import stirling.software.SPDF.model.User;
@Component
 public class CustomAuthenticationFailureHandler extends SimpleUrlAuthenticationFailureHandler {

-    @Autowired private final LoginAttemptService loginAttemptService;
+    private LoginAttemptService loginAttemptService;

-    @Autowired private final UserService userService; // Inject the UserService
+    private UserService userService;
+
+    private static final Logger logger =
+            LoggerFactory.getLogger(CustomAuthenticationFailureHandler.class);

    public CustomAuthenticationFailureHandler(
            LoginAttemptService loginAttemptService, UserService userService) {
@@ -34,22 +39,28 @@ public class CustomAuthenticationFailureHandler extends SimpleUrlAuthenticationF
            HttpServletResponse response,
            AuthenticationException exception)
            throws IOException, ServletException {
+
        String ip = request.getRemoteAddr();
        logger.error("Failed login attempt from IP: " + ip);

        String username = request.getParameter("username");
        if (!isDemoUser(username)) {
            if (loginAttemptService.loginAttemptCheck(username)) {
-                setDefaultFailureUrl("/login?error=locked");
-
+                response.sendRedirect("/login?error=locked");
+                return;
            } else {
                if (exception.getClass().isAssignableFrom(LockedException.class)) {
-                    setDefaultFailureUrl("/login?error=locked");
+                    response.sendRedirect("/login?error=locked");
+                    return;
+                } else if (exception instanceof UsernameNotFoundException) {
+                    response.sendRedirect("/login?error=oauth2AuthenticationError");
+                    return;
                }
            }
        }
        if (exception.getClass().isAssignableFrom(BadCredentialsException.class)) {
-            setDefaultFailureUrl("/login?error=badcredentials");
+            response.sendRedirect("/login?error=badcredentials");
+            return;
        }

        super.onAuthenticationFailure(request, response, exception);
--- a/src/main/java/stirling/software/SPDF/config/security/CustomAuthenticationSuccessHandler.java
+++ b/src/main/java/stirling/software/SPDF/config/security/CustomAuthenticationSuccessHandler.java
@@ -2,7 +2,6 @@ package stirling.software.SPDF.config.security;

 import java.io.IOException;

-import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.security.core.Authentication;
 import org.springframework.security.web.authentication.SavedRequestAwareAuthenticationSuccessHandler;
 import org.springframework.security.web.savedrequest.SavedRequest;
@@ -18,15 +17,10 @@ import stirling.software.SPDF.utils.RequestUriUtils;
 public class CustomAuthenticationSuccessHandler
        extends SavedRequestAwareAuthenticationSuccessHandler {

-    @Autowired private LoginAttemptService loginAttemptService;
-
    @Override
    public void onAuthenticationSuccess(
            HttpServletRequest request, HttpServletResponse response, Authentication authentication)
            throws ServletException, IOException {
-        String username = request.getParameter("username");
-        loginAttemptService.loginSucceeded(username);
-
        // Get the saved request
        HttpSession session = request.getSession(false);
        SavedRequest savedRequest =
--- a/src/main/java/stirling/software/SPDF/config/security/CustomLogoutSuccessHandler.java
+++ b/src/main/java/stirling/software/SPDF/config/security/CustomLogoutSuccessHandler.java
@@ -2,10 +2,9 @@ package stirling.software.SPDF.config.security;

 import java.io.IOException;

-import org.springframework.context.annotation.Bean;
+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.security.core.Authentication;
 import org.springframework.security.core.session.SessionRegistry;
-import org.springframework.security.core.session.SessionRegistryImpl;
 import org.springframework.security.web.authentication.logout.SimpleUrlLogoutSuccessHandler;

 import jakarta.servlet.ServletException;
@@ -14,10 +13,8 @@ import jakarta.servlet.http.HttpServletResponse;
 import jakarta.servlet.http.HttpSession;

 public class CustomLogoutSuccessHandler extends SimpleUrlLogoutSuccessHandler {
-    @Bean
-    public SessionRegistry sessionRegistry() {
-        return new SessionRegistryImpl();
-    }
+
+    @Autowired SessionRegistry sessionRegistry;

    @Override
    public void onLogoutSuccess(
@@ -26,14 +23,9 @@ public class CustomLogoutSuccessHandler extends SimpleUrlLogoutSuccessHandler {
        HttpSession session = request.getSession(false);
        if (session != null) {
            String sessionId = session.getId();
-            sessionRegistry().removeSessionInformation(sessionId);
+            sessionRegistry.removeSessionInformation(sessionId);
        }

-        if (request.getParameter("oauth2AutoCreateDisabled") != null) {
-            response.sendRedirect(
-                    request.getContextPath() + "/login?error=oauth2AutoCreateDisabled");
-        } else {
-            response.sendRedirect(request.getContextPath() + "/login?logout=true");
-        }
+        response.sendRedirect(request.getContextPath() + "/login?logout=true");
    }
 }
--- a/src/main/java/stirling/software/SPDF/config/security/CustomUserDetailsService.java
+++ b/src/main/java/stirling/software/SPDF/config/security/CustomUserDetailsService.java
@@ -22,7 +22,11 @@ public class CustomUserDetailsService implements UserDetailsService {

    @Autowired private UserRepository userRepository;

-    @Autowired private LoginAttemptService loginAttemptService;
+    private LoginAttemptService loginAttemptService;
+
+    CustomUserDetailsService(LoginAttemptService loginAttemptService) {
+        this.loginAttemptService = loginAttemptService;
+    }

    @Override
    public UserDetails loadUserByUsername(String username) throws UsernameNotFoundException {
@@ -39,6 +43,10 @@ public class CustomUserDetailsService implements UserDetailsService {
                    "Your account has been locked due to too many failed login attempts.");
        }

+        if (user.getPassword() == null || user.getPassword().isEmpty()) {
+            throw new UsernameNotFoundException("Password must not be null");
+        }
+
        return new org.springframework.security.core.userdetails.User(
                user.getUsername(),
                user.getPassword(),
--- a/src/main/java/stirling/software/SPDF/config/security/InitialSecuritySetup.java
+++ b/src/main/java/stirling/software/SPDF/config/security/InitialSecuritySetup.java
@@ -7,6 +7,8 @@ import java.nio.file.Paths;
 import java.util.List;
 import java.util.UUID;

+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.stereotype.Component;

@@ -21,6 +23,8 @@ public class InitialSecuritySetup {

    @Autowired ApplicationProperties applicationProperties;

+    private static final Logger logger = LoggerFactory.getLogger(InitialSecuritySetup.class);
+
    @PostConstruct
    public void init() {
        if (!userService.hasUsers()) {
@@ -29,6 +33,20 @@ public class InitialSecuritySetup {
                    applicationProperties.getSecurity().getInitialLogin().getUsername();
            String initialPassword =
                    applicationProperties.getSecurity().getInitialLogin().getPassword();
+            try {
+                // https://github.com/Stirling-Tools/Stirling-PDF/issues/976
+                userService.isUsernameValidWithReturn(initialUsername);
+            } catch (IllegalArgumentException e) {
+                Path pathToFile = Paths.get("configs/settings.yml");
+                try {
+                    if (Files.exists(pathToFile)) {
+                        Files.delete(pathToFile);
+                    }
+                } catch (IOException ex) {
+                    logger.info(ex.getMessage());
+                }
+                throw e;
+            }
            if (initialUsername != null && initialPassword != null) {
                userService.saveUser(initialUsername, initialPassword, Role.ADMIN.getRoleId());
            } else {
--- a/src/main/java/stirling/software/SPDF/config/security/LoginAttemptService.java
+++ b/src/main/java/stirling/software/SPDF/config/security/LoginAttemptService.java
@@ -30,28 +30,31 @@ public class LoginAttemptService {
            new ConcurrentHashMap<>();

    public void loginSucceeded(String key) {
-        attemptsCache.remove(key);
+        attemptsCache.remove(key.toLowerCase());
    }

    public boolean loginAttemptCheck(String key) {
-        attemptsCache.compute(
-                key,
-                (k, attemptCounter) -> {
-                    if (attemptCounter == null
-                            || attemptCounter.shouldReset(ATTEMPT_INCREMENT_TIME)) {
-                        return new AttemptCounter();
-                    } else {
-                        attemptCounter.increment();
-                        return attemptCounter;
-                    }
-                });
-        return attemptsCache.get(key).getAttemptCount() >= MAX_ATTEMPTS;
+        return attemptsCache
+                        .compute(
+                                key.toLowerCase(),
+                                (k, attemptCounter) -> {
+                                    if (attemptCounter == null
+                                            || attemptCounter.shouldReset(ATTEMPT_INCREMENT_TIME)) {
+                                        return new AttemptCounter();
+                                    } else {
+                                        attemptCounter.increment();
+                                        return attemptCounter;
+                                    }
+                                })
+                        .getAttemptCount()
+                >= MAX_ATTEMPTS;
    }

    public boolean isBlocked(String key) {
-        AttemptCounter attemptCounter = attemptsCache.get(key);
+        AttemptCounter attemptCounter = attemptsCache.get(key.toLowerCase());
        if (attemptCounter != null) {
-            return attemptCounter.getAttemptCount() >= MAX_ATTEMPTS;
+            return attemptCounter.getAttemptCount() >= MAX_ATTEMPTS
+                    && !attemptCounter.shouldReset(ATTEMPT_INCREMENT_TIME);
        }
        return false;
    }
--- a/src/main/java/stirling/software/SPDF/config/security/SecurityConfiguration.java
+++ b/src/main/java/stirling/software/SPDF/config/security/SecurityConfiguration.java
@@ -1,6 +1,5 @@
 package stirling.software.SPDF.config.security;

-import java.io.IOException;
 import java.util.*;

 import org.springframework.beans.factory.annotation.Autowired;
@@ -14,33 +13,30 @@ import org.springframework.security.config.annotation.method.configuration.Enabl
 import org.springframework.security.config.annotation.web.builders.HttpSecurity;
 import org.springframework.security.config.annotation.web.configuration.EnableWebSecurity;
 import org.springframework.security.config.http.SessionCreationPolicy;
-import org.springframework.security.core.Authentication;
 import org.springframework.security.core.GrantedAuthority;
 import org.springframework.security.core.authority.SimpleGrantedAuthority;
 import org.springframework.security.core.authority.mapping.GrantedAuthoritiesMapper;
 import org.springframework.security.core.session.SessionRegistry;
 import org.springframework.security.core.session.SessionRegistryImpl;
-import org.springframework.security.core.userdetails.UserDetailsService;
 import org.springframework.security.crypto.bcrypt.BCryptPasswordEncoder;
 import org.springframework.security.crypto.password.PasswordEncoder;
 import org.springframework.security.oauth2.client.registration.ClientRegistration;
 import org.springframework.security.oauth2.client.registration.ClientRegistrationRepository;
 import org.springframework.security.oauth2.client.registration.ClientRegistrations;
 import org.springframework.security.oauth2.client.registration.InMemoryClientRegistrationRepository;
-import org.springframework.security.oauth2.core.user.OAuth2User;
 import org.springframework.security.oauth2.core.user.OAuth2UserAuthority;
 import org.springframework.security.web.SecurityFilterChain;
-import org.springframework.security.web.authentication.AuthenticationSuccessHandler;
 import org.springframework.security.web.authentication.UsernamePasswordAuthenticationFilter;
 import org.springframework.security.web.authentication.rememberme.PersistentTokenRepository;
 import org.springframework.security.web.savedrequest.NullRequestCache;
 import org.springframework.security.web.util.matcher.AntPathRequestMatcher;

-import jakarta.servlet.ServletException;
-import jakarta.servlet.http.HttpServletRequest;
-import jakarta.servlet.http.HttpServletResponse;
-import jakarta.servlet.http.HttpSession;
+import stirling.software.SPDF.config.security.oauth2.CustomOAuth2AuthenticationFailureHandler;
+import stirling.software.SPDF.config.security.oauth2.CustomOAuth2AuthenticationSuccessHandler;
+import stirling.software.SPDF.config.security.oauth2.CustomOAuth2LogoutSuccessHandler;
+import stirling.software.SPDF.config.security.oauth2.CustomOAuth2UserService;
 import stirling.software.SPDF.model.ApplicationProperties;
+import stirling.software.SPDF.model.ApplicationProperties.Security.OAUTH2;
 import stirling.software.SPDF.model.User;
 import stirling.software.SPDF.repository.JPATokenRepositoryImpl;

@@ -49,7 +45,7 @@ import stirling.software.SPDF.repository.JPATokenRepositoryImpl;
@EnableMethodSecurity
 public class SecurityConfiguration {

-    @Autowired private UserDetailsService userDetailsService;
+    @Autowired private CustomUserDetailsService userDetailsService;

    @Bean
    public PasswordEncoder passwordEncoder() {
@@ -111,18 +107,7 @@ public class SecurityConfiguration {
                                                    new AntPathRequestMatcher("/logout"))
                                            .logoutSuccessHandler(new CustomLogoutSuccessHandler())
                                            .invalidateHttpSession(true) // Invalidate session
-                                            .deleteCookies("JSESSIONID", "remember-me")
-                                            .addLogoutHandler(
-                                                    (request, response, authentication) -> {
-                                                        HttpSession session =
-                                                                request.getSession(false);
-                                                        if (session != null) {
-                                                            String sessionId = session.getId();
-                                                            sessionRegistry()
-                                                                    .removeSessionInformation(
-                                                                            sessionId);
-                                                        }
-                                                    }))
+                                            .deleteCookies("JSESSIONID", "remember-me"))
                    .rememberMe(
                            rememberMeConfigurer ->
                                    rememberMeConfigurer // Use the configurator directly
@@ -168,43 +153,33 @@ public class SecurityConfiguration {
            if (applicationProperties.getSecurity().getOAUTH2().getEnabled()) {

                http.oauth2Login(
-                        oauth2 ->
-                                oauth2.loginPage("/oauth2")
-                                        /*
-                                        This Custom handler is used to check if the OAUTH2 user trying to log in, already exists in the database.
-                                        If user exists, login proceeds as usual. If user does not exist, then it is autocreated but only if 'OAUTH2AutoCreateUser'
-                                        is set as true, else login fails with an error message advising the same.
-                                         */
-                                        .successHandler(
-                                                new AuthenticationSuccessHandler() {
-                                                    @Override
-                                                    public void onAuthenticationSuccess(
-                                                            HttpServletRequest request,
-                                                            HttpServletResponse response,
-                                                            Authentication authentication)
-                                                            throws ServletException, IOException {
-                                                        OAuth2User oauthUser =
-                                                                (OAuth2User)
-                                                                        authentication
-                                                                                .getPrincipal();
-                                                        if (userService.processOAuth2PostLogin(
-                                                                oauthUser.getAttribute("email"),
-                                                                applicationProperties
-                                                                        .getSecurity()
-                                                                        .getOAUTH2()
-                                                                        .getAutoCreateUser())) {
-                                                            response.sendRedirect("/");
-                                                        } else {
-                                                            response.sendRedirect(
-                                                                    "/logout?oauth2AutoCreateDisabled=true");
-                                                        }
-                                                    }
-                                                })
-                                        // Add existing Authorities from the database
-                                        .userInfoEndpoint(
-                                                userInfoEndpoint ->
-                                                        userInfoEndpoint.userAuthoritiesMapper(
-                                                                userAuthoritiesMapper())));
+                                oauth2 ->
+                                        oauth2.loginPage("/oauth2")
+                                                /*
+                                                This Custom handler is used to check if the OAUTH2 user trying to log in, already exists in the database.
+                                                If user exists, login proceeds as usual. If user does not exist, then it is autocreated but only if 'OAUTH2AutoCreateUser'
+                                                is set as true, else login fails with an error message advising the same.
+                                                 */
+                                                .successHandler(
+                                                        new CustomOAuth2AuthenticationSuccessHandler(
+                                                                applicationProperties, userService))
+                                                .failureHandler(
+                                                        new CustomOAuth2AuthenticationFailureHandler())
+                                                // Add existing Authorities from the database
+                                                .userInfoEndpoint(
+                                                        userInfoEndpoint ->
+                                                                userInfoEndpoint
+                                                                        .oidcUserService(
+                                                                                new CustomOAuth2UserService(
+                                                                                        applicationProperties))
+                                                                        .userAuthoritiesMapper(
+                                                                                userAuthoritiesMapper())))
+                        .userDetailsService(userDetailsService)
+                        .logout(
+                                logout ->
+                                        logout.logoutSuccessHandler(
+                                                new CustomOAuth2LogoutSuccessHandler(
+                                                        this.applicationProperties)));
            }
        } else {
            http.csrf(csrf -> csrf.disable())
@@ -225,13 +200,13 @@ public class SecurityConfiguration {
    }

    private ClientRegistration oidcClientRegistration() {
-        return ClientRegistrations.fromOidcIssuerLocation(
-                        applicationProperties.getSecurity().getOAUTH2().getIssuer())
+        OAUTH2 oauth = applicationProperties.getSecurity().getOAUTH2();
+        return ClientRegistrations.fromIssuerLocation(oauth.getIssuer())
                .registrationId("oidc")
-                .clientId(applicationProperties.getSecurity().getOAUTH2().getClientId())
-                .clientSecret(applicationProperties.getSecurity().getOAUTH2().getClientSecret())
-                .scope("openid", "profile", "email")
-                .userNameAttributeName("email")
+                .clientId(oauth.getClientId())
+                .clientSecret(oauth.getClientSecret())
+                .scope(oauth.getScopes())
+                .userNameAttributeName(oauth.getUseAsUsername())
                .clientName("OIDC")
                .build();
    }
@@ -256,9 +231,14 @@ public class SecurityConfiguration {

                        // Add Authorities from database for existing user, if user is present.
                        if (authority instanceof OAuth2UserAuthority oauth2Auth) {
+                            String useAsUsername =
+                                    applicationProperties
+                                            .getSecurity()
+                                            .getOAUTH2()
+                                            .getUseAsUsername();
                            Optional<User> userOpt =
                                    userService.findByUsernameIgnoreCase(
-                                            (String) oauth2Auth.getAttributes().get("email"));
+                                            (String) oauth2Auth.getAttributes().get(useAsUsername));
                            if (userOpt.isPresent()) {
                                User user = userOpt.get();
                                if (user != null) {
--- a/src/main/java/stirling/software/SPDF/config/security/UserService.java
+++ b/src/main/java/stirling/software/SPDF/config/security/UserService.java
@@ -8,6 +8,8 @@ import java.util.UUID;
 import java.util.stream.Collectors;

 import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.context.MessageSource;
+import org.springframework.context.i18n.LocaleContextHolder;
 import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
 import org.springframework.security.core.Authentication;
 import org.springframework.security.core.GrantedAuthority;
@@ -18,6 +20,7 @@ import org.springframework.security.crypto.password.PasswordEncoder;
 import org.springframework.stereotype.Service;

 import stirling.software.SPDF.controller.api.pipeline.UserServiceInterface;
+import stirling.software.SPDF.model.AuthenticationType;
 import stirling.software.SPDF.model.Authority;
 import stirling.software.SPDF.model.Role;
 import stirling.software.SPDF.model.User;
@@ -33,19 +36,19 @@ public class UserService implements UserServiceInterface {

    @Autowired private PasswordEncoder passwordEncoder;

+    @Autowired private MessageSource messageSource;
+
    // Handle OAUTH2 login and user auto creation.
    public boolean processOAuth2PostLogin(String username, boolean autoCreateUser) {
+        if (!isUsernameValidWithReturn(username).equals(username)) {
+            return false;
+        }
        Optional<User> existUser = userRepository.findByUsernameIgnoreCase(username);
        if (existUser.isPresent()) {
            return true;
        }
        if (autoCreateUser) {
-            User user = new User();
-            user.setUsername(username);
-            user.setEnabled(true);
-            user.setFirstLogin(false);
-            user.addAuthority(new Authority(Role.USER.getRoleId(), user));
-            userRepository.save(user);
+            saveUser(username, AuthenticationType.OAUTH2);
            return true;
        }
        return false;
@@ -128,30 +131,46 @@ public class UserService implements UserServiceInterface {
        return userOpt.isPresent() && userOpt.get().getApiKey().equals(apiKey);
    }

-    public void saveUser(String username, String password) {
+    public void saveUser(String username, AuthenticationType authenticationType)
+            throws IllegalArgumentException {
        User user = new User();
-        user.setUsername(username);
-        user.setPassword(passwordEncoder.encode(password));
+        user.setUsername(isUsernameValidWithReturn(username));
        user.setEnabled(true);
+        user.setFirstLogin(false);
+        user.addAuthority(new Authority(Role.USER.getRoleId(), user));
+        user.setAuthenticationType(authenticationType);
        userRepository.save(user);
    }

-    public void saveUser(String username, String password, String role, boolean firstLogin) {
+    public void saveUser(String username, String password) throws IllegalArgumentException {
        User user = new User();
-        user.setUsername(username);
+        user.setUsername(isUsernameValidWithReturn(username));
+        user.setPassword(passwordEncoder.encode(password));
+        user.setEnabled(true);
+        user.setAuthenticationType(AuthenticationType.WEB);
+        userRepository.save(user);
+    }
+
+    public void saveUser(String username, String password, String role, boolean firstLogin)
+            throws IllegalArgumentException {
+        User user = new User();
+        user.setUsername(isUsernameValidWithReturn(username));
        user.setPassword(passwordEncoder.encode(password));
        user.addAuthority(new Authority(role, user));
        user.setEnabled(true);
+        user.setAuthenticationType(AuthenticationType.WEB);
        user.setFirstLogin(firstLogin);
        userRepository.save(user);
    }

-    public void saveUser(String username, String password, String role) {
+    public void saveUser(String username, String password, String role)
+            throws IllegalArgumentException {
        User user = new User();
-        user.setUsername(username);
+        user.setUsername(isUsernameValidWithReturn(username));
        user.setPassword(passwordEncoder.encode(password));
        user.addAuthority(new Authority(role, user));
        user.setEnabled(true);
+        user.setAuthenticationType(AuthenticationType.WEB);
        user.setFirstLogin(false);
        userRepository.save(user);
    }
@@ -209,8 +228,8 @@ public class UserService implements UserServiceInterface {
        return authorityRepository.findByUserId(user.getId());
    }

-    public void changeUsername(User user, String newUsername) {
-        user.setUsername(newUsername);
+    public void changeUsername(User user, String newUsername) throws IllegalArgumentException {
+        user.setUsername(isUsernameValidWithReturn(newUsername));
        userRepository.save(user);
    }

@@ -235,6 +254,40 @@ public class UserService implements UserServiceInterface {
    }

    public boolean isUsernameValid(String username) {
-        return username.matches("[a-zA-Z0-9]+");
+        // Checks whether the simple username is formatted correctly
+        boolean isValidSimpleUsername =
+                username.matches("^[a-zA-Z0-9][a-zA-Z0-9@._+-]*[a-zA-Z0-9]$");
+        // Checks whether the email address is formatted correctly
+        boolean isValidEmail =
+                username.matches(
+                        "^(?=.{1,64}@)[A-Za-z0-9]+(\\.[A-Za-z0-9_+.-]+)*@[^-][A-Za-z0-9-]+(\\.[A-Za-z0-9-]+)*(\\.[A-Za-z]{2,})$");
+        return isValidSimpleUsername || isValidEmail;
+    }
+
+    public String isUsernameValidWithReturn(String username) throws IllegalArgumentException {
+        if (!isUsernameValid(username)) {
+            String message =
+                    messageSource.getMessage(
+                            "invalidUsernameMessage", null, LocaleContextHolder.getLocale());
+            throw new IllegalArgumentException(message);
+        }
+        return username;
+    }
+
+    public boolean hasPassword(String username) {
+        Optional<User> user = userRepository.findByUsernameIgnoreCase(username);
+        if (user.isPresent() && user.get().hasPassword()) {
+            return true;
+        }
+        return false;
+    }
+
+    public boolean isAuthenticationTypeByUsername(
+            String username, AuthenticationType authenticationType) {
+        Optional<User> user = userRepository.findByUsernameIgnoreCase(username);
+        if (user.isPresent() && user.get().getAuthenticationType() != null) {
+            return user.get().getAuthenticationType().equalsIgnoreCase(authenticationType.name());
+        }
+        return false;
    }
 }
--- a/src/main/java/stirling/software/SPDF/config/security/oauth2/CustomOAuth2AuthenticationFailureHandler.java
+++ b/src/main/java/stirling/software/SPDF/config/security/oauth2/CustomOAuth2AuthenticationFailureHandler.java
@@ -0,0 +1,40 @@
+package stirling.software.SPDF.config.security.oauth2;
+
+import java.io.IOException;
+
+import org.springframework.security.authentication.LockedException;
+import org.springframework.security.core.AuthenticationException;
+import org.springframework.security.core.userdetails.UsernameNotFoundException;
+import org.springframework.security.oauth2.core.OAuth2AuthenticationException;
+import org.springframework.security.oauth2.core.OAuth2Error;
+import org.springframework.security.web.authentication.SimpleUrlAuthenticationFailureHandler;
+import org.springframework.stereotype.Component;
+
+import jakarta.servlet.ServletException;
+import jakarta.servlet.http.HttpServletRequest;
+import jakarta.servlet.http.HttpServletResponse;
+
+@Component
+public class CustomOAuth2AuthenticationFailureHandler
+        extends SimpleUrlAuthenticationFailureHandler {
+
+    @Override
+    public void onAuthenticationFailure(
+            HttpServletRequest request,
+            HttpServletResponse response,
+            AuthenticationException exception)
+            throws IOException, ServletException {
+        if (exception instanceof OAuth2AuthenticationException) {
+            OAuth2Error error = ((OAuth2AuthenticationException) exception).getError();
+            getRedirectStrategy()
+                    .sendRedirect(request, response, "/login?error=oAuth::" + error.getErrorCode());
+        } else if (exception instanceof LockedException) {
+            getRedirectStrategy().sendRedirect(request, response, "/login?error=locked");
+        } else if (exception instanceof UsernameNotFoundException) {
+            getRedirectStrategy()
+                    .sendRedirect(request, response, "/login?error=oauth2AuthenticationError");
+        } else {
+            super.onAuthenticationFailure(request, response, exception);
+        }
+    }
+}
--- a/src/main/java/stirling/software/SPDF/config/security/oauth2/CustomOAuth2AuthenticationSuccessHandler.java
+++ b/src/main/java/stirling/software/SPDF/config/security/oauth2/CustomOAuth2AuthenticationSuccessHandler.java
@@ -0,0 +1,74 @@
+package stirling.software.SPDF.config.security.oauth2;
+
+import java.io.IOException;
+
+import org.springframework.security.core.Authentication;
+import org.springframework.security.oauth2.core.user.OAuth2User;
+import org.springframework.security.web.authentication.SavedRequestAwareAuthenticationSuccessHandler;
+import org.springframework.security.web.savedrequest.SavedRequest;
+import org.springframework.stereotype.Component;
+
+import jakarta.servlet.ServletException;
+import jakarta.servlet.http.HttpServletRequest;
+import jakarta.servlet.http.HttpServletResponse;
+import jakarta.servlet.http.HttpSession;
+import stirling.software.SPDF.config.security.UserService;
+import stirling.software.SPDF.model.ApplicationProperties;
+import stirling.software.SPDF.model.ApplicationProperties.Security.OAUTH2;
+import stirling.software.SPDF.model.AuthenticationType;
+import stirling.software.SPDF.utils.RequestUriUtils;
+
+@Component
+public class CustomOAuth2AuthenticationSuccessHandler
+        extends SavedRequestAwareAuthenticationSuccessHandler {
+
+    ApplicationProperties applicationProperties;
+    UserService userService;
+
+    public CustomOAuth2AuthenticationSuccessHandler(
+            ApplicationProperties applicationProperties, UserService userService) {
+        this.applicationProperties = applicationProperties;
+        this.userService = userService;
+    }
+
+    @Override
+    public void onAuthenticationSuccess(
+            HttpServletRequest request, HttpServletResponse response, Authentication authentication)
+            throws ServletException, IOException {
+
+        OAuth2User oauthUser = (OAuth2User) authentication.getPrincipal();
+
+        // Get the saved request
+        HttpSession session = request.getSession(false);
+        SavedRequest savedRequest =
+                session != null
+                        ? (SavedRequest) session.getAttribute("SPRING_SECURITY_SAVED_REQUEST")
+                        : null;
+        if (savedRequest != null
+                && !RequestUriUtils.isStaticResource(savedRequest.getRedirectUrl())) {
+            // Redirect to the original destination
+            super.onAuthenticationSuccess(request, response, authentication);
+        } else {
+            OAUTH2 oAuth = applicationProperties.getSecurity().getOAUTH2();
+            String username = oauthUser.getAttribute(oAuth.getUseAsUsername());
+            if (userService.usernameExistsIgnoreCase(username)
+                    && userService.hasPassword(username)
+                    && !userService.isAuthenticationTypeByUsername(
+                            username, AuthenticationType.OAUTH2)
+                    && oAuth.getAutoCreateUser()) {
+                response.sendRedirect(
+                        request.getContextPath() + "/logout?oauth2AuthenticationError=true");
+                return;
+            } else {
+                try {
+                    userService.processOAuth2PostLogin(username, oAuth.getAutoCreateUser());
+                    response.sendRedirect("/");
+                    return;
+                } catch (IllegalArgumentException e) {
+                    response.sendRedirect("/logout?invalidUsername=true");
+                    return;
+                }
+            }
+        }
+    }
+}
--- a/src/main/java/stirling/software/SPDF/config/security/oauth2/CustomOAuth2LogoutSuccessHandler.java
+++ b/src/main/java/stirling/software/SPDF/config/security/oauth2/CustomOAuth2LogoutSuccessHandler.java
@@ -0,0 +1,87 @@
+package stirling.software.SPDF.config.security.oauth2;
+
+import java.io.IOException;
+
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.security.core.Authentication;
+import org.springframework.security.core.session.SessionRegistry;
+import org.springframework.security.core.userdetails.UserDetails;
+import org.springframework.security.oauth2.core.user.OAuth2User;
+import org.springframework.security.web.authentication.logout.SimpleUrlLogoutSuccessHandler;
+import org.springframework.stereotype.Component;
+
+import jakarta.servlet.ServletException;
+import jakarta.servlet.http.HttpServletRequest;
+import jakarta.servlet.http.HttpServletResponse;
+import jakarta.servlet.http.HttpSession;
+import stirling.software.SPDF.model.ApplicationProperties;
+import stirling.software.SPDF.model.ApplicationProperties.Security.OAUTH2;
+
+@Component
+public class CustomOAuth2LogoutSuccessHandler extends SimpleUrlLogoutSuccessHandler {
+
+    @Autowired SessionRegistry sessionRegistry;
+
+    private ApplicationProperties applicationProperties;
+
+    public CustomOAuth2LogoutSuccessHandler(ApplicationProperties applicationProperties) {
+        this.applicationProperties = applicationProperties;
+    }
+
+    @Override
+    public void onLogoutSuccess(
+            HttpServletRequest request, HttpServletResponse response, Authentication authentication)
+            throws IOException, ServletException {
+
+        boolean isOAuthUser = true;
+        String param = "logout=true";
+        if (authentication == null) {
+            response.sendRedirect("/");
+            return;
+        }
+        Object pri = authentication.getPrincipal();
+        if (pri instanceof UserDetails) {
+            UserDetails userDetails = (UserDetails) pri;
+            isOAuthUser = userDetails.getPassword() == null;
+        } else if (pri instanceof OAuth2User) {
+            isOAuthUser = true;
+        }
+
+        OAUTH2 oauth = applicationProperties.getSecurity().getOAUTH2();
+        String provider = oauth.getProvider() != null && isOAuthUser ? oauth.getProvider() : "";
+
+        if (request.getParameter("oauth2AuthenticationError") != null) {
+            param = "error=oauth2AuthenticationError";
+        } else if (request.getParameter("invalidUsername") != null) {
+            param = "error=invalidUsername";
+        }
+        HttpSession session = request.getSession(false);
+        if (session != null) {
+            String sessionId = session.getId();
+            sessionRegistry.removeSessionInformation(sessionId);
+        }
+
+        switch (provider) {
+            case "keycloak":
+                String logoutUrl =
+                        oauth.getIssuer()
+                                + "/protocol/openid-connect/logout"
+                                + "?client_id="
+                                + oauth.getClientId()
+                                + "&post_logout_redirect_uri="
+                                + response.encodeRedirectURL(
+                                        "http://" + request.getHeader("host") + "/login?" + param);
+                response.sendRedirect(logoutUrl);
+                break;
+            case "google":
+            default:
+                if (request.getParameter("oauth2AutoCreateDisabled") != null) {
+                    response.sendRedirect(
+                            request.getContextPath() + "/login?error=oauth2AutoCreateDisabled");
+                } else {
+                    response.sendRedirect(request.getContextPath() + "/login?logout=true");
+                }
+                break;
+        }
+    }
+}
--- a/src/main/java/stirling/software/SPDF/config/security/oauth2/CustomOAuth2UserService.java
+++ b/src/main/java/stirling/software/SPDF/config/security/oauth2/CustomOAuth2UserService.java
@@ -0,0 +1,52 @@
+package stirling.software.SPDF.config.security.oauth2;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import org.springframework.security.oauth2.client.oidc.userinfo.OidcUserRequest;
+import org.springframework.security.oauth2.client.oidc.userinfo.OidcUserService;
+import org.springframework.security.oauth2.client.userinfo.OAuth2UserService;
+import org.springframework.security.oauth2.core.OAuth2AuthenticationException;
+import org.springframework.security.oauth2.core.OAuth2Error;
+import org.springframework.security.oauth2.core.oidc.user.DefaultOidcUser;
+import org.springframework.security.oauth2.core.oidc.user.OidcUser;
+
+import stirling.software.SPDF.model.ApplicationProperties;
+
+public class CustomOAuth2UserService implements OAuth2UserService<OidcUserRequest, OidcUser> {
+
+    private final OidcUserService delegate = new OidcUserService();
+
+    private ApplicationProperties applicationProperties;
+
+    public CustomOAuth2UserService(ApplicationProperties applicationProperties) {
+        this.applicationProperties = applicationProperties;
+    }
+
+    @Override
+    public OidcUser loadUser(OidcUserRequest userRequest) throws OAuth2AuthenticationException {
+        String usernameAttribute =
+                applicationProperties.getSecurity().getOAUTH2().getUseAsUsername();
+        try {
+
+            OidcUser user = delegate.loadUser(userRequest);
+            Map<String, Object> attributes = new HashMap<>(user.getAttributes());
+
+            // Ensure the preferred username attribute is present
+            if (!attributes.containsKey(usernameAttribute)) {
+                attributes.put(usernameAttribute, attributes.getOrDefault("email", ""));
+                usernameAttribute = "email";
+            }
+
+            // Return a new OidcUser with adjusted attributes
+            return new DefaultOidcUser(
+                    user.getAuthorities(),
+                    userRequest.getIdToken(),
+                    user.getUserInfo(),
+                    usernameAttribute);
+        } catch (java.lang.IllegalArgumentException e) {
+            throw new OAuth2AuthenticationException(
+                    new OAuth2Error(e.getMessage()), e.getMessage(), e);
+        }
+    }
+}
--- a/src/main/java/stirling/software/SPDF/config/security/oauth2/CustomOAuthUserService.java
+++ b/src/main/java/stirling/software/SPDF/config/security/oauth2/CustomOAuthUserService.java
@@ -0,0 +1,57 @@
+package stirling.software.SPDF.config.security.oauth2;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.security.oauth2.client.oidc.userinfo.OidcUserRequest;
+import org.springframework.security.oauth2.client.oidc.userinfo.OidcUserService;
+import org.springframework.security.oauth2.client.userinfo.OAuth2UserService;
+import org.springframework.security.oauth2.core.OAuth2AuthenticationException;
+import org.springframework.security.oauth2.core.OAuth2Error;
+import org.springframework.security.oauth2.core.oidc.user.DefaultOidcUser;
+import org.springframework.security.oauth2.core.oidc.user.OidcUser;
+
+import stirling.software.SPDF.model.ApplicationProperties;
+
+public class CustomOAuthUserService implements OAuth2UserService<OidcUserRequest, OidcUser> {
+
+    private static final Logger logger = LoggerFactory.getLogger(CustomOAuthUserService.class);
+
+    private final OidcUserService delegate = new OidcUserService();
+
+    private ApplicationProperties applicationProperties;
+
+    public CustomOAuthUserService(ApplicationProperties applicationProperties) {
+        this.applicationProperties = applicationProperties;
+    }
+
+    @Override
+    public OidcUser loadUser(OidcUserRequest userRequest) throws OAuth2AuthenticationException {
+        String usernameAttribute =
+                applicationProperties.getSecurity().getOAUTH2().getUseAsUsername();
+        try {
+
+            OidcUser user = delegate.loadUser(userRequest);
+            Map<String, Object> attributes = new HashMap<>(user.getAttributes());
+
+            // Ensure the preferred username attribute is present
+            if (!attributes.containsKey(usernameAttribute)) {
+                attributes.put(usernameAttribute, attributes.getOrDefault("email", ""));
+                usernameAttribute = "email";
+                logger.info("Adjusted username attribute to use email");
+            }
+
+            // Return a new OidcUser with adjusted attributes
+            return new DefaultOidcUser(
+                    user.getAuthorities(),
+                    userRequest.getIdToken(),
+                    user.getUserInfo(),
+                    usernameAttribute);
+        } catch (java.lang.IllegalArgumentException e) {
+            throw new OAuth2AuthenticationException(
+                    new OAuth2Error(e.getMessage()), e.getMessage(), e);
+        }
+    }
+}