Key Rotation

/**
 * Interface for locating public keys used in token verification with support for key rotation
 */
public interface PublicKeyLocator {
    /**
     * Get public key for token verification
     * @param kid Key identifier from token header, may be null
     * @param deployment Keycloak deployment configuration
     * @return PublicKey for verification or null if not found
     */
    PublicKey getPublicKey(String kid, KeycloakDeployment deployment);
    
    /**
     * Reset/clear cached keys, forcing fresh retrieval
     * @param deployment Keycloak deployment configuration
     */
    void reset(KeycloakDeployment deployment);
}

/**
 * Public key locator implementation using a single hardcoded public key
 */
public class HardcodedPublicKeyLocator implements PublicKeyLocator {
    /**
     * Constructor with hardcoded public key
     * @param publicKey Public key to use for all token verification
     */
    public HardcodedPublicKeyLocator(PublicKey publicKey);
    
    /**
     * Get the hardcoded public key
     * @param kid Key identifier (ignored)
     * @param deployment Keycloak deployment (ignored)
     * @return The configured public key
     */
    public PublicKey getPublicKey(String kid, KeycloakDeployment deployment);
    
    /**
     * Reset operation (no-op for hardcoded key)
     * @param deployment Keycloak deployment (ignored)
     */
    public void reset(KeycloakDeployment deployment);
}

// Create hardcoded key locator from certificate
Certificate cert = loadCertificateFromFile("/path/to/realm-certificate.pem");
PublicKey publicKey = cert.getPublicKey();
PublicKeyLocator keyLocator = new HardcodedPublicKeyLocator(publicKey);

// Configure in deployment
KeycloakDeployment deployment = new KeycloakDeployment();
deployment.setPublicKeyLocator(keyLocator);

// Use for token verification
AccessToken token = AdapterTokenVerifier.verifyToken(tokenString, deployment);

/**
 * Public key locator that downloads keys from Keycloak's JWKS endpoint with caching and rotation support
 */
public class JWKPublicKeyLocator implements PublicKeyLocator {
    /**
     * Default constructor
     */
    public JWKPublicKeyLocator();
    
    /**
     * Get public key from JWKS endpoint, with caching
     * @param kid Key identifier from token header
     * @param deployment Keycloak deployment configuration
     * @return PublicKey for verification or null if not found
     */
    public PublicKey getPublicKey(String kid, KeycloakDeployment deployment);
    
    /**
     * Reset cached keys, forcing fresh download from JWKS endpoint
     * @param deployment Keycloak deployment configuration
     */
    public void reset(KeycloakDeployment deployment);
}

// Create JWKS-based key locator (recommended for production)
PublicKeyLocator keyLocator = new JWKPublicKeyLocator();

// Configure deployment with JWKS support
KeycloakDeployment deployment = new KeycloakDeployment();
deployment.setPublicKeyLocator(keyLocator);
deployment.setMinTimeBetweenJwksRequests(10); // Minimum 10 seconds between JWKS requests
deployment.setPublicKeyCacheTtl(3600); // Cache keys for 1 hour

// The locator will automatically:
// 1. Download JWKS from deployment.getJwksUrl()
// 2. Cache keys with TTL
// 3. Handle key rotation automatically
// 4. Respect rate limiting for JWKS requests

// Token verification with automatic key resolution
try {
    AccessToken token = AdapterTokenVerifier.verifyToken(tokenString, deployment);
    // Token verified successfully with appropriate key
} catch (VerificationException e) {
    // Verification failed - could be due to key rotation
    // Reset keys and retry once
    keyLocator.reset(deployment);
    try {
        AccessToken token = AdapterTokenVerifier.verifyToken(tokenString, deployment);
    } catch (VerificationException retryEx) {
        // Verification failed after key refresh
        logger.warn("Token verification failed: {}", retryEx.getMessage());
    }
}

/**
 * Utility class for token verification with integrated public key resolution
 */
public class AdapterTokenVerifier {
    /**
     * Container for verified tokens
     */
    public static class VerifiedTokens {
        // Implementation details - contains verified access and ID tokens
    }
    
    /**
     * Verify access token using deployment configuration
     * @param tokenString Raw JWT token string
     * @param deployment Keycloak deployment with key locator
     * @return Verified AccessToken instance
     * @throws VerificationException If token verification fails
     */
    public static AccessToken verifyToken(String tokenString, KeycloakDeployment deployment) 
        throws VerificationException;
    
    /**
     * Verify both access token and ID token
     * @param accessTokenString Raw access token JWT string
     * @param idTokenString Raw ID token JWT string
     * @param deployment Keycloak deployment with key locator
     * @return VerifiedTokens container with both verified tokens
     * @throws VerificationException If either token verification fails
     */
    public static VerifiedTokens verifyTokens(
        String accessTokenString,
        String idTokenString,
        KeycloakDeployment deployment
    ) throws VerificationException;
    
    /**
     * Create token verifier with custom configuration
     * @param tokenString Raw JWT token string
     * @param deployment Keycloak deployment configuration
     * @param withDefaultChecks Whether to include default verification checks
     * @param tokenClass Class of token to verify (AccessToken, IDToken, etc.)
     * @return Configured TokenVerifier instance
     * @throws VerificationException If verifier creation fails
     */
    public static <T extends JsonWebToken> TokenVerifier<T> createVerifier(
        String tokenString,
        KeycloakDeployment deployment,
        boolean withDefaultChecks,
        Class<T> tokenClass
    ) throws VerificationException;
}

// Basic token verification
try {
    AccessToken accessToken = AdapterTokenVerifier.verifyToken(tokenString, deployment);
    
    // Token is valid, extract information
    String username = accessToken.getPreferredUsername();
    String email = accessToken.getEmail();
    Set<String> roles = accessToken.getRealmAccess().getRoles();
    long expiration = accessToken.getExpiration();
    
    // Check if token is still valid (not expired)
    if (expiration < System.currentTimeMillis() / 1000) {
        logger.warn("Token has expired");
    }
    
} catch (VerificationException e) {
    logger.warn("Token verification failed: {}", e.getMessage());
    // Handle invalid token (return 401, redirect to login, etc.)
}

// Verify both access and ID tokens
try {
    AdapterTokenVerifier.VerifiedTokens tokens = AdapterTokenVerifier.verifyTokens(
        accessTokenString, 
        idTokenString, 
        deployment
    );
    
    // Both tokens verified successfully
    // Access token information and ID token information are available
    
} catch (VerificationException e) {
    logger.warn("Token verification failed: {}", e.getMessage());
}

// Custom token verification with specific checks
try {
    TokenVerifier<AccessToken> verifier = AdapterTokenVerifier.createVerifier(
        tokenString,
        deployment,
        true, // Include default checks (signature, expiration, audience, etc.)
        AccessToken.class
    );
    
    // Add custom verification checks
    verifier.check(token -> {
        // Custom validation logic
        if (!token.getIssuer().equals(expectedIssuer)) {
            throw new VerificationException("Invalid issuer");
        }
        
        // Check custom claims
        Object customClaim = token.getOtherClaims().get("custom_claim");
        if (customClaim == null) {
            throw new VerificationException("Missing required custom claim");
        }
    });
    
    AccessToken token = verifier.verify();
    
} catch (VerificationException e) {
    logger.warn("Custom token verification failed: {}", e.getMessage());
}

// Database-backed public key locator with caching
public class DatabasePublicKeyLocator implements PublicKeyLocator {
    private final PublicKeyRepository keyRepository;
    private final Map<String, PublicKey> keyCache = new ConcurrentHashMap<>();
    private final Map<String, Long> cacheTimestamps = new ConcurrentHashMap<>();
    private final long cacheTtlMs = TimeUnit.HOURS.toMillis(1);
    
    public DatabasePublicKeyLocator(PublicKeyRepository keyRepository) {
        this.keyRepository = keyRepository;
    }
    
    @Override
    public PublicKey getPublicKey(String kid, KeycloakDeployment deployment) {
        String realm = deployment.getRealm();
        String cacheKey = realm + ":" + (kid != null ? kid : "default");
        
        // Check cache first
        PublicKey cachedKey = keyCache.get(cacheKey);
        Long cacheTime = cacheTimestamps.get(cacheKey);
        
        if (cachedKey != null && cacheTime != null && 
            (System.currentTimeMillis() - cacheTime) < cacheTtlMs) {
            return cachedKey;
        }
        
        // Load from database
        PublicKeyEntity keyEntity = keyRepository.findByRealmAndKid(realm, kid);
        if (keyEntity != null) {
            try {
                PublicKey publicKey = parsePublicKey(keyEntity.getKeyData());
                keyCache.put(cacheKey, publicKey);
                cacheTimestamps.put(cacheKey, System.currentTimeMillis());
                return publicKey;
            } catch (Exception e) {
                logger.warn("Failed to parse public key from database", e);
            }
        }
        
        return null;
    }
    
    @Override
    public void reset(KeycloakDeployment deployment) {
        String realm = deployment.getRealm();
        // Remove cached keys for this realm
        keyCache.entrySet().removeIf(entry -> entry.getKey().startsWith(realm + ":"));
        cacheTimestamps.entrySet().removeIf(entry -> entry.getKey().startsWith(realm + ":"));
    }
    
    private PublicKey parsePublicKey(String keyData) throws Exception {
        // Parse PEM or other format
        // Implementation depends on key format
    }
}

// Composite key locator with fallback strategy
public class CompositePublicKeyLocator implements PublicKeyLocator {
    private final List<PublicKeyLocator> locators;
    
    public CompositePublicKeyLocator(PublicKeyLocator... locators) {
        this.locators = Arrays.asList(locators);
    }
    
    @Override
    public PublicKey getPublicKey(String kid, KeycloakDeployment deployment) {
        for (PublicKeyLocator locator : locators) {
            try {
                PublicKey key = locator.getPublicKey(kid, deployment);
                if (key != null) {
                    return key;
                }
            } catch (Exception e) {
                logger.debug("Key locator {} failed: {}", locator.getClass().getSimpleName(), e.getMessage());
            }
        }
        return null;
    }
    
    @Override
    public void reset(KeycloakDeployment deployment) {
        for (PublicKeyLocator locator : locators) {
            try {
                locator.reset(deployment);
            } catch (Exception e) {
                logger.warn("Failed to reset key locator {}", locator.getClass().getSimpleName(), e);
            }
        }
    }
}

// Usage: JWKS first, then hardcoded fallback
PublicKeyLocator composite = new CompositePublicKeyLocator(
    new JWKPublicKeyLocator(),
    new HardcodedPublicKeyLocator(fallbackPublicKey)
);

// Monitor key rotation events
public class KeyRotationMonitor {
    private final KeycloakDeployment deployment;
    private final PublicKeyLocator keyLocator;
    private volatile String currentKeyId;
    
    public KeyRotationMonitor(KeycloakDeployment deployment) {
        this.deployment = deployment;
        this.keyLocator = deployment.getPublicKeyLocator();
    }
    
    public boolean handleTokenVerificationFailure(String tokenString, VerificationException e) {
        try {
            // Parse token header to get key ID
            String kid = extractKeyIdFromToken(tokenString);
            
            if (kid != null && !kid.equals(currentKeyId)) {
                logger.info("Detected potential key rotation. Current: {}, Token: {}", currentKeyId, kid);
                
                // Reset key cache and retry
                keyLocator.reset(deployment);
                
                // Attempt verification with fresh keys
                AccessToken token = AdapterTokenVerifier.verifyToken(tokenString, deployment);
                currentKeyId = kid;
                
                logger.info("Token verification successful after key rotation");
                return true;
            }
        } catch (Exception retryException) {
            logger.warn("Token verification failed even after key rotation handling", retryException);
        }
        
        return false;
    }
    
    private String extractKeyIdFromToken(String tokenString) {
        try {
            String[] parts = tokenString.split("\\.");
            if (parts.length >= 2) {
                String headerJson = new String(Base64.getUrlDecoder().decode(parts[0]));
                // Parse JSON to extract "kid" field
                // Implementation depends on JSON library
            }
        } catch (Exception e) {
            logger.debug("Failed to extract key ID from token", e);
        }
        return null;
    }
}

// Production configuration with JWKS and caching
KeycloakDeployment deployment = new KeycloakDeployment();
deployment.setRealm("production-realm");
deployment.setAuthServerBaseUrl("https://keycloak.company.com/auth");
deployment.setPublicKeyLocator(new JWKPublicKeyLocator());
deployment.setMinTimeBetweenJwksRequests(10); // Rate limit JWKS requests
deployment.setPublicKeyCacheTtl(3600); // Cache for 1 hour

// Development configuration with hardcoded key
PublicKey devKey = loadDevelopmentPublicKey();
KeycloakDeployment devDeployment = new KeycloakDeployment();
devDeployment.setPublicKeyLocator(new HardcodedPublicKeyLocator(devKey));

// High-availability configuration with multiple strategies
PublicKeyLocator haLocator = new CompositePublicKeyLocator(
    new JWKPublicKeyLocator(),                    // Primary: JWKS
    new DatabasePublicKeyLocator(keyRepository),  // Fallback: Database
    new HardcodedPublicKeyLocator(emergencyKey)   // Emergency: Hardcoded
);
deployment.setPublicKeyLocator(haLocator);