Token Storage

/**
 * Interface defining the contract for token storage implementations (per-request object)
 */
public interface AdapterTokenStore {
    /**
     * Check the current token validity and refresh if necessary
     */
    void checkCurrentToken();
    
    /**
     * Check if authenticator results are cached
     * @param authenticator The request authenticator to check
     * @return true if cached, false otherwise
     */
    boolean isCached(RequestAuthenticator authenticator);
    
    /**
     * Save account information after successful authentication
     * @param account Authenticated account to store
     */
    void saveAccountInfo(OidcKeycloakAccount account);
    
    /**
     * Logout and clear stored token information
     */
    void logout();
    
    /**
     * Callback invoked after token refresh
     * @param securityContext Updated security context with new tokens
     */
    void refreshCallback(RefreshableKeycloakSecurityContext securityContext);
}

// Custom token store implementation
public class DatabaseTokenStore implements AdapterTokenStore {
    private final String sessionId;
    private final TokenRepository tokenRepository;
    
    public DatabaseTokenStore(String sessionId, TokenRepository repository) {
        this.sessionId = sessionId;
        this.tokenRepository = repository;
    }
    
    @Override
    public void checkCurrentToken() {
        TokenInfo tokenInfo = tokenRepository.findBySessionId(sessionId);
        if (tokenInfo != null && tokenInfo.isExpired()) {
            // Token expired, trigger refresh
            refreshToken(tokenInfo);
        }
    }
    
    @Override
    public boolean isCached(RequestAuthenticator authenticator) {
        return tokenRepository.existsBySessionId(sessionId);
    }
    
    @Override
    public void saveAccountInfo(OidcKeycloakAccount account) {
        RefreshableKeycloakSecurityContext context = 
            (RefreshableKeycloakSecurityContext) account.getKeycloakSecurityContext();
        
        TokenInfo tokenInfo = new TokenInfo();
        tokenInfo.setSessionId(sessionId);
        tokenInfo.setAccessToken(context.getTokenString());
        tokenInfo.setRefreshToken(context.getRefreshToken());
        tokenInfo.setIdToken(context.getIdTokenString());
        tokenInfo.setExpiresAt(context.getToken().getExpiration());
        
        tokenRepository.save(tokenInfo);
    }
    
    @Override
    public void logout() {
        tokenRepository.deleteBySessionId(sessionId);
    }
    
    @Override
    public void refreshCallback(RefreshableKeycloakSecurityContext securityContext) {
        // Update stored tokens after refresh
        TokenInfo tokenInfo = tokenRepository.findBySessionId(sessionId);
        if (tokenInfo != null) {
            tokenInfo.setAccessToken(securityContext.getTokenString());
            tokenInfo.setRefreshToken(securityContext.getRefreshToken());
            tokenInfo.setExpiresAt(securityContext.getToken().getExpiration());
            tokenRepository.save(tokenInfo);
        }
    }
}

/**
 * Utility class for cookie-based token storage operations
 */
public class CookieTokenStore {
    /**
     * Set authentication token as a secure cookie
     * @param deployment Keycloak deployment configuration
     * @param facade HTTP facade for cookie operations
     * @param session Security context containing token information
     */
    public static void setTokenCookie(
        KeycloakDeployment deployment,
        HttpFacade facade,
        RefreshableKeycloakSecurityContext session
    );
    
    /**
     * Retrieve principal from authentication cookie
     * @param deployment Keycloak deployment configuration
     * @param facade HTTP facade for cookie operations
     * @param tokenStore Token store for caching
     * @return KeycloakPrincipal if valid cookie found, null otherwise
     */
    public static KeycloakPrincipal<RefreshableKeycloakSecurityContext> getPrincipalFromCookie(
        KeycloakDeployment deployment,
        HttpFacade facade,
        AdapterTokenStore tokenStore
    );
    
    /**
     * Remove authentication cookie
     * @param deployment Keycloak deployment configuration  
     * @param facade HTTP facade for cookie operations
     */
    public static void removeCookie(KeycloakDeployment deployment, HttpFacade facade);
}

// Set authentication cookie after successful login
RefreshableKeycloakSecurityContext securityContext = 
    (RefreshableKeycloakSecurityContext) principal.getKeycloakSecurityContext();
CookieTokenStore.setTokenCookie(deployment, httpFacade, securityContext);

// Retrieve principal from cookie on subsequent requests
KeycloakPrincipal<RefreshableKeycloakSecurityContext> principal = 
    CookieTokenStore.getPrincipalFromCookie(deployment, httpFacade, tokenStore);

if (principal != null) {
    // User authenticated via cookie
    RefreshableKeycloakSecurityContext context = principal.getKeycloakSecurityContext();
    AccessToken token = context.getToken();
    
    // Check if token needs refresh
    if (!context.isTokenTimeToLiveSufficient(token)) {
        boolean refreshed = context.refreshExpiredToken(true);
        if (refreshed) {
            // Update cookie with new tokens
            CookieTokenStore.setTokenCookie(deployment, httpFacade, context);
        } else {
            // Refresh failed, remove invalid cookie
            CookieTokenStore.removeCookie(deployment, httpFacade);
        }
    }
}

// Logout - remove cookie
CookieTokenStore.removeCookie(deployment, httpFacade);

// HTTP session-based token store
public class SessionTokenStore implements AdapterTokenStore {
    private final HttpSession session;
    private static final String ACCOUNT_KEY = "keycloak.account";
    
    public SessionTokenStore(HttpSession session) {
        this.session = session;
    }
    
    @Override
    public void saveAccountInfo(OidcKeycloakAccount account) {
        session.setAttribute(ACCOUNT_KEY, account);
    }
    
    @Override
    public boolean isCached(RequestAuthenticator authenticator) {
        return session.getAttribute(ACCOUNT_KEY) != null;
    }
    
    @Override
    public void logout() {
        session.removeAttribute(ACCOUNT_KEY);
        session.invalidate();
    }
    
    public OidcKeycloakAccount getAccount() {
        return (OidcKeycloakAccount) session.getAttribute(ACCOUNT_KEY);
    }
}

// Redis-based distributed token store
public class RedisTokenStore implements AdapterTokenStore {
    private final RedisTemplate<String, Object> redisTemplate;
    private final String sessionId;
    private final Duration tokenTtl;
    
    public RedisTokenStore(RedisTemplate<String, Object> redisTemplate, 
                          String sessionId, Duration tokenTtl) {
        this.redisTemplate = redisTemplate;
        this.sessionId = sessionId;
        this.tokenTtl = tokenTtl;
    }
    
    @Override
    public void saveAccountInfo(OidcKeycloakAccount account) {
        String key = "keycloak:session:" + sessionId;
        RefreshableKeycloakSecurityContext context = 
            (RefreshableKeycloakSecurityContext) account.getKeycloakSecurityContext();
        
        Map<String, Object> tokenData = new HashMap<>();
        tokenData.put("accessToken", context.getTokenString());
        tokenData.put("refreshToken", context.getRefreshToken());
        tokenData.put("idToken", context.getIdTokenString());
        tokenData.put("principal", account.getPrincipal());
        tokenData.put("roles", account.getRoles());
        
        redisTemplate.opsForHash().putAll(key, tokenData);
        redisTemplate.expire(key, tokenTtl);
    }
    
    @Override
    public boolean isCached(RequestAuthenticator authenticator) {
        String key = "keycloak:session:" + sessionId;
        return redisTemplate.hasKey(key);
    }
    
    @Override
    public void logout() {
        String key = "keycloak:session:" + sessionId;
        redisTemplate.delete(key);
    }
    
    @Override
    public void refreshCallback(RefreshableKeycloakSecurityContext securityContext) {
        // Update tokens in Redis after refresh
        String key = "keycloak:session:" + sessionId;
        if (redisTemplate.hasKey(key)) {
            redisTemplate.opsForHash().put(key, "accessToken", securityContext.getTokenString());
            redisTemplate.opsForHash().put(key, "refreshToken", securityContext.getRefreshToken());
            redisTemplate.opsForHash().put(key, "idToken", securityContext.getIdTokenString());
            redisTemplate.expire(key, tokenTtl);
        }
    }
}

// Custom JWT-based cookie storage
public class JwtCookieTokenStore implements AdapterTokenStore {
    private final HttpFacade facade;
    private final KeycloakDeployment deployment;
    private final String cookieName;
    
    public JwtCookieTokenStore(HttpFacade facade, KeycloakDeployment deployment) {
        this.facade = facade;
        this.deployment = deployment;
        this.cookieName = deployment.getStateCookieName() + "_token";
    }
    
    @Override
    public void saveAccountInfo(OidcKeycloakAccount account) {
        RefreshableKeycloakSecurityContext context = 
            (RefreshableKeycloakSecurityContext) account.getKeycloakSecurityContext();
        
        // Create compact token info for cookie
        Map<String, Object> tokenInfo = new HashMap<>();
        tokenInfo.put("sub", context.getToken().getSubject());
        tokenInfo.put("exp", context.getToken().getExpiration());
        tokenInfo.put("roles", account.getRoles());
        tokenInfo.put("rt", context.getRefreshToken()); // Refresh token
        
        // Sign and serialize as JWT
        String jwtToken = createSignedJwt(tokenInfo);
        
        // Set secure cookie
        Cookie cookie = new Cookie(cookieName, jwtToken);
        cookie.setPath(deployment.getAdapterStateCookiePath());
        cookie.setSecure(deployment.getSslRequired() == SslRequired.ALL);
        cookie.setHttpOnly(true);
        cookie.setMaxAge((int) (context.getToken().getExpiration() - System.currentTimeMillis() / 1000));
        
        facade.getResponse().setCookie(cookie);
    }
    
    @Override
    public boolean isCached(RequestAuthenticator authenticator) {
        Cookie cookie = facade.getRequest().getCookie(cookieName);
        return cookie != null && isValidJwtCookie(cookie.getValue());
    }
    
    @Override
    public void logout() {
        Cookie cookie = new Cookie(cookieName, "");
        cookie.setPath(deployment.getAdapterStateCookiePath());
        cookie.setMaxAge(0);
        facade.getResponse().setCookie(cookie);
    }
    
    private String createSignedJwt(Map<String, Object> claims) {
        // Implementation depends on JWT library (e.g., jose4j, nimbus-jose-jwt)
        // Sign with deployment's client secret or private key
        return "signed.jwt.token";
    }
    
    private boolean isValidJwtCookie(String jwt) {
        try {
            // Verify JWT signature and expiration
            Map<String, Object> claims = parseAndVerifyJwt(jwt);
            long exp = (Long) claims.get("exp");
            return exp > System.currentTimeMillis() / 1000;
        } catch (Exception e) {
            return false;
        }
    }
}

// Configure token store in deployment
KeycloakDeployment deployment = new KeycloakDeployment();
deployment.setTokenStore(TokenStore.COOKIE); // or TokenStore.SESSION

// Configure cookie settings
deployment.setAdapterStateCookiePath("/app");
deployment.setStateCookieName("KEYCLOAK_ADAPTER");

// For custom storage implementations
public class CustomAdapterTokenStoreFactory {
    public static AdapterTokenStore create(HttpFacade facade, KeycloakDeployment deployment) {
        TokenStore tokenStore = deployment.getTokenStore();
        
        switch (tokenStore) {
            case COOKIE:
                return new CookieBasedTokenStore(facade, deployment);
            case SESSION:
                return new SessionTokenStore(facade.getRequest().getSession());
            default:
                // Custom implementation
                return new RedisTokenStore(getRedisTemplate(), getSessionId(facade), Duration.ofHours(1));
        }
    }
}