JAAS Integration

/**
 * Base abstract login module providing common JAAS integration functionality for Keycloak authentication
 */
public abstract class AbstractKeycloakLoginModule implements LoginModule {
    /**
     * Configuration option key for Keycloak configuration file path
     */
    public static final String KEYCLOAK_CONFIG_FILE_OPTION = "keycloak-config-file";
    
    /**
     * Configuration option key for role principal class name
     */
    public static final String ROLE_PRINCIPAL_CLASS_OPTION = "role-principal-class";
    
    /**
     * Resource path for profile configuration
     */
    public static final String PROFILE_RESOURCE = "profile-resource";
    
    /**
     * Authentication result container
     */
    public static class Auth {
        // Contains authentication details and tokens
    }
    
    /**
     * Initialize the login module with configuration
     * @param subject Subject to authenticate
     * @param callbackHandler Handler for obtaining credentials
     * @param sharedState Shared state between login modules
     * @param options Configuration options
     */
    public void initialize(
        Subject subject,
        CallbackHandler callbackHandler,
        Map<String, ?> sharedState,
        Map<String, ?> options
    );
    
    /**
     * Perform authentication
     * @return true if authentication succeeded
     * @throws LoginException If authentication fails
     */
    public boolean login() throws LoginException;
    
    /**
     * Commit the authentication (add principals to subject)
     * @return true if commit succeeded
     * @throws LoginException If commit fails
     */
    public boolean commit() throws LoginException;
    
    /**
     * Abort the authentication (cleanup on failure)
     * @return true if abort succeeded
     * @throws LoginException If abort fails
     */
    public boolean abort() throws LoginException;
    
    /**
     * Logout (remove principals from subject)
     * @return true if logout succeeded
     * @throws LoginException If logout fails
     */
    public boolean logout() throws LoginException;
    
    /**
     * Resolve Keycloak deployment from configuration file
     * @param keycloakConfigFile Path to Keycloak configuration file
     * @return Resolved KeycloakDeployment
     * @throws RuntimeException If configuration cannot be loaded
     */
    protected KeycloakDeployment resolveDeployment(String keycloakConfigFile);
    
    /**
     * Create role principal for the given role name
     * @param roleName Name of the role
     * @return Principal representing the role
     */
    protected Principal createRolePrincipal(String roleName);
    
    /**
     * Authenticate using bearer token
     * @param tokenString Bearer token string
     * @return Authentication result
     * @throws VerificationException If token verification fails
     */
    protected Auth bearerAuth(String tokenString) throws VerificationException;
    
    /**
     * Post-process authentication result after token verification
     * @param tokenString Original token string
     * @param token Verified access token
     * @return Updated authentication result
     */
    protected Auth postTokenVerification(String tokenString, AccessToken token);
    
    /**
     * Perform authentication with username and password (abstract method)
     * @param username Username credential
     * @param password Password credential
     * @return Authentication result
     * @throws Exception If authentication fails
     */
    protected abstract Auth doAuth(String username, String password) throws Exception;
    
    /**
     * Get logger for this login module (abstract method)
     * @return Logger instance
     */
    protected abstract Logger getLogger();
}

// Custom login module implementation
public class CustomKeycloakLoginModule extends AbstractKeycloakLoginModule {
    private static final Logger logger = LoggerFactory.getLogger(CustomKeycloakLoginModule.class);
    
    @Override
    protected Auth doAuth(String username, String password) throws Exception {
        // Implement username/password authentication
        KeycloakDeployment deployment = resolveDeployment(getConfigFile());
        
        // Use Keycloak's direct access grants (resource owner password credentials)
        try {
            AccessTokenResponse tokenResponse = authenticateWithPassword(deployment, username, password);
            String tokenString = tokenResponse.getToken();
            
            // Verify the token
            AccessToken token = AdapterTokenVerifier.verifyToken(tokenString, deployment);
            
            // Return authentication result
            return postTokenVerification(tokenString, token);
            
        } catch (Exception e) {
            getLogger().warn("Authentication failed for user: {}", username, e);
            throw new LoginException("Authentication failed: " + e.getMessage());
        }
    }
    
    @Override
    protected Logger getLogger() {
        return logger;
    }
    
    private AccessTokenResponse authenticateWithPassword(KeycloakDeployment deployment, 
                                                       String username, String password) throws Exception {
        // Implementation for direct access grant
        HttpClient client = deployment.getClient();
        HttpPost post = new HttpPost(deployment.getTokenUrl());
        
        List<NameValuePair> formParams = new ArrayList<>();
        formParams.add(new BasicNameValuePair("grant_type", "password"));
        formParams.add(new BasicNameValuePair("username", username));
        formParams.add(new BasicNameValuePair("password", password));
        
        AdapterUtils.setClientCredentials(deployment, post, formParams);
        
        post.setEntity(new UrlEncodedFormEntity(formParams));
        
        HttpResponse response = client.execute(post);
        int statusCode = response.getStatusLine().getStatusCode();
        
        if (statusCode == 200) {
            return JsonSerialization.readValue(response.getEntity().getContent(), AccessTokenResponse.class);
        } else {
            throw new Exception("Authentication failed with status: " + statusCode);
        }
    }
}

// JAAS configuration file (jaas.conf)
/*
MyApplication {
    com.example.CustomKeycloakLoginModule required
        keycloak-config-file="/path/to/keycloak.json"
        role-principal-class="org.keycloak.adapters.jaas.RolePrincipal";
};
*/

// Usage in application
System.setProperty("java.security.auth.login.config", "/path/to/jaas.conf");

LoginContext loginContext = new LoginContext("MyApplication", new CallbackHandler() {
    @Override
    public void handle(Callback[] callbacks) throws IOException, UnsupportedCallbackException {
        for (Callback callback : callbacks) {
            if (callback instanceof NameCallback) {
                ((NameCallback) callback).setName("user@example.com");
            } else if (callback instanceof PasswordCallback) {
                ((PasswordCallback) callback).setPassword("password".toCharArray());
            }
        }
    }
});

try {
    loginContext.login();
    Subject subject = loginContext.getSubject();
    
    // Access authenticated user information
    Set<Principal> principals = subject.getPrincipals();
    for (Principal principal : principals) {
        System.out.println("Principal: " + principal.getName());
    }
    
    // Perform authorized actions
    Subject.doAs(subject, new PrivilegedAction<Void>() {
        @Override
        public Void run() {
            // Code running with authenticated subject
            return null;
        }
    });
    
} finally {
    loginContext.logout();
}

/**
 * Principal implementation representing a user role in JAAS context
 */
public class RolePrincipal implements Principal {
    /**
     * Constructor with role name
     * @param roleName Name of the role
     */
    public RolePrincipal(String roleName);
    
    /**
     * Check equality with another principal
     * @param p Principal to compare with
     * @return true if principals are equal
     */
    public boolean equals(Object p);
    
    /**
     * Get hash code for this principal
     * @return Hash code value
     */
    public int hashCode();
    
    /**
     * Get the role name
     * @return Role name string
     */
    public String getName();
    
    /**
     * String representation of the role principal
     * @return Formatted string representation
     */
    public String toString();
}

// Create role principals
RolePrincipal adminRole = new RolePrincipal("admin");
RolePrincipal userRole = new RolePrincipal("user");

// Use in custom login module
@Override
public boolean commit() throws LoginException {
    if (authenticationSucceeded) {
        // Add user principal
        subject.getPrincipals().add(new KeycloakPrincipal<>(username, securityContext));
        
        // Add role principals
        Set<String> roles = extractRolesFromToken(accessToken);
        for (String role : roles) {
            subject.getPrincipals().add(new RolePrincipal(role));
        }
        
        return true;
    }
    return false;
}

// Check roles in authorized code
public boolean hasRole(Subject subject, String roleName) {
    Set<RolePrincipal> rolePrincipals = subject.getPrincipals(RolePrincipal.class);
    return rolePrincipals.stream()
        .anyMatch(role -> roleName.equals(role.getName()));
}

// Authorization check example
Subject.doAs(subject, new PrivilegedAction<Void>() {
    @Override
    public Void run() {
        Subject currentSubject = Subject.getSubject(AccessController.getContext());
        
        if (hasRole(currentSubject, "admin")) {
            // Perform admin operations
            performAdminOperation();
        } else if (hasRole(currentSubject, "user")) {
            // Perform user operations
            performUserOperation();
        } else {
            throw new SecurityException("Insufficient privileges");
        }
        
        return null;
    }
});

// Login module for bearer token authentication
public class BearerTokenLoginModule extends AbstractKeycloakLoginModule {
    private static final Logger logger = LoggerFactory.getLogger(BearerTokenLoginModule.class);
    private String bearerToken;
    
    @Override
    public void initialize(Subject subject, CallbackHandler callbackHandler, 
                          Map<String, ?> sharedState, Map<String, ?> options) {
        super.initialize(subject, callbackHandler, sharedState, options);
        
        // Extract bearer token from shared state or callback
        this.bearerToken = (String) sharedState.get("bearer.token");
    }
    
    @Override
    protected Auth doAuth(String username, String password) throws Exception {
        // This module doesn't use username/password
        throw new UnsupportedOperationException("Use bearer token authentication");
    }
    
    @Override
    public boolean login() throws LoginException {
        if (bearerToken == null) {
            // Try to get token via callback
            try {
                BearerTokenCallback tokenCallback = new BearerTokenCallback();
                callbackHandler.handle(new Callback[]{tokenCallback});
                bearerToken = tokenCallback.getToken();
            } catch (Exception e) {
                throw new LoginException("Failed to obtain bearer token: " + e.getMessage());
            }
        }
        
        if (bearerToken == null) {
            return false;
        }
        
        try {
            auth = bearerAuth(bearerToken);
            return auth != null;
        } catch (VerificationException e) {
            getLogger().warn("Bearer token verification failed", e);
            throw new LoginException("Invalid bearer token: " + e.getMessage());
        }
    }
    
    @Override
    protected Logger getLogger() {
        return logger;
    }
}

// Custom callback for bearer token
public class BearerTokenCallback implements Callback {
    private String token;
    
    public String getToken() {
        return token;
    }
    
    public void setToken(String token) {
        this.token = token;
    }
}

// Servlet filter integrating JAAS with web requests
public class JAASKeycloakFilter implements Filter {
    private String jaasConfigName;
    
    @Override
    public void init(FilterConfig filterConfig) throws ServletException {
        jaasConfigName = filterConfig.getInitParameter("jaas-config-name");
        if (jaasConfigName == null) {
            jaasConfigName = "KeycloakWeb";
        }
    }
    
    @Override
    public void doFilter(ServletRequest request, ServletResponse response, 
                        FilterChain chain) throws IOException, ServletException {
        HttpServletRequest httpRequest = (HttpServletRequest) request;
        HttpServletResponse httpResponse = (HttpServletResponse) response;
        
        String authHeader = httpRequest.getHeader("Authorization");
        if (authHeader != null && authHeader.startsWith("Bearer ")) {
            String token = authHeader.substring(7);
            
            try {
                // Authenticate using JAAS
                LoginContext loginContext = new LoginContext(jaasConfigName, new BearerTokenCallbackHandler(token));
                loginContext.login();
                
                Subject subject = loginContext.getSubject();
                
                // Continue with authenticated subject
                Subject.doAs(subject, new PrivilegedExceptionAction<Void>() {
                    @Override
                    public Void run() throws Exception {
                        chain.doFilter(request, response);
                        return null;
                    }
                });
                
                loginContext.logout();
                
            } catch (LoginException e) {
                httpResponse.setStatus(HttpServletResponse.SC_UNAUTHORIZED);
                httpResponse.getWriter().write("Authentication failed: " + e.getMessage());
                return;
            } catch (PrivilegedActionException e) {
                throw new ServletException(e.getException());
            }
        } else {
            httpResponse.setStatus(HttpServletResponse.SC_UNAUTHORIZED);
            httpResponse.getWriter().write("Bearer token required");
        }
    }
    
    private static class BearerTokenCallbackHandler implements CallbackHandler {
        private final String token;
        
        public BearerTokenCallbackHandler(String token) {
            this.token = token;
        }
        
        @Override
        public void handle(Callback[] callbacks) throws IOException, UnsupportedCallbackException {
            for (Callback callback : callbacks) {
                if (callback instanceof BearerTokenCallback) {
                    ((BearerTokenCallback) callback).setToken(token);
                } else {
                    throw new UnsupportedCallbackException(callback);
                }
            }
        }
    }
}

// JBoss/WildFly integration example
public class KeycloakSecurityDomain {
    
    // Configure in standalone.xml or domain.xml
    /*
    <security-domain name="keycloak">
        <authentication>
            <login-module code="com.example.KeycloakLoginModule" flag="required">
                <module-option name="keycloak-config-file" value="/opt/keycloak/keycloak.json"/>
                <module-option name="role-principal-class" value="org.keycloak.adapters.jaas.RolePrincipal"/>
            </login-module>
        </authentication>
    </security-domain>
    */
    
    // EJB with security annotations
    @Stateless
    @SecurityDomain("keycloak")
    public class SecureEJB {
        
        @RolesAllowed({"admin", "manager"})
        public void adminOperation() {
            Subject subject = Subject.getSubject(AccessController.getContext());
            // Access authenticated subject
        }
        
        @RolesAllowed("user")
        public void userOperation() {
            // User-level operation
        }
    }
}

// Complete JAAS configuration
/*
# jaas.conf
KeycloakApp {
    com.example.CustomKeycloakLoginModule required
        keycloak-config-file="/etc/keycloak/keycloak.json"
        role-principal-class="org.keycloak.adapters.jaas.RolePrincipal";
};

KeycloakWeb {
    com.example.BearerTokenLoginModule required
        keycloak-config-file="/etc/keycloak/keycloak.json";
};
*/

// System properties setup
System.setProperty("java.security.auth.login.config", "/etc/jaas.conf");

// Programmatic configuration
Configuration.setConfiguration(new Configuration() {
    @Override
    public AppConfigurationEntry[] getAppConfigurationEntry(String name) {
        if ("KeycloakApp".equals(name)) {
            Map<String, String> options = new HashMap<>();
            options.put("keycloak-config-file", "/etc/keycloak/keycloak.json");
            options.put("role-principal-class", "org.keycloak.adapters.jaas.RolePrincipal");
            
            return new AppConfigurationEntry[] {
                new AppConfigurationEntry(
                    "com.example.CustomKeycloakLoginModule",
                    AppConfigurationEntry.LoginModuleControlFlag.REQUIRED,
                    options
                )
            };
        }
        return null;
    }
});

/**
 * JAAS login module for bearer token authentication
 * Expects username (ignored) and password (bearer token)
 */
public class BearerTokenLoginModule extends AbstractKeycloakLoginModule {
    /**
     * Authenticate using bearer token passed as password parameter
     * @param username Username (ignored for bearer token auth)
     * @param password Bearer token string
     * @return Authentication result
     * @throws VerificationException If token verification fails
     */
    protected Auth doAuth(String username, String password) throws VerificationException;
    
    /**
     * Get logger for this login module
     * @return Logger instance
     */
    protected Logger getLogger();
}

/**
 * JAAS login module for direct access grants (username/password authentication)
 * Implements OAuth2 Resource Owner Password Credentials Grant
 */
public class DirectAccessGrantsLoginModule extends AbstractKeycloakLoginModule {
    /**
     * Configuration option key for OAuth2 scope parameter
     */
    public static final String SCOPE_OPTION = "scope";
    
    /**
     * Authenticate using direct access grants (username/password)
     * @param username User's username
     * @param password User's password
     * @return Authentication result with tokens
     * @throws IOException If network communication fails
     * @throws VerificationException If token verification fails
     */
    protected Auth doAuth(String username, String password) throws IOException, VerificationException;
    
    /**
     * Perform direct grant authentication against Keycloak
     * @param username User's username
     * @param password User's password
     * @return Authentication result with access and refresh tokens
     * @throws IOException If HTTP request fails
     * @throws VerificationException If token verification fails
     */
    protected Auth directGrantAuth(String username, String password) throws IOException, VerificationException;
    
    /**
     * Commit authentication (save refresh token to subject's private credentials)
     * @return true if commit succeeded
     * @throws LoginException If commit fails
     */
    public boolean commit() throws LoginException;
    
    /**
     * Logout and revoke refresh token with Keycloak
     * @return true if logout succeeded
     * @throws LoginException If logout fails
     */
    public boolean logout() throws LoginException;
}