PoC for a CustomKeycloakOIDCFilter with PKCE support

CustomKeycloakOIDCFilterWithPkceSupport.java

package demo;

import org.keycloak.adapters.AdapterTokenStore;
import org.keycloak.adapters.AuthenticatedActionsHandler;
import org.keycloak.adapters.KeycloakDeployment;
import org.keycloak.adapters.OAuthRequestAuthenticator;
import org.keycloak.adapters.OIDCAuthenticationError;
import org.keycloak.adapters.OIDCHttpFacade;
import org.keycloak.adapters.PreAuthActionsHandler;
import org.keycloak.adapters.RequestAuthenticator;
import org.keycloak.adapters.ServerRequest;
import org.keycloak.adapters.rotation.AdapterTokenVerifier;
import org.keycloak.adapters.servlet.FilterRequestAuthenticator;
import org.keycloak.adapters.servlet.KeycloakOIDCFilter;
import org.keycloak.adapters.servlet.OIDCFilterSessionStore;
import org.keycloak.adapters.servlet.OIDCServletHttpFacade;
import org.keycloak.adapters.spi.AdapterSessionStore;
import org.keycloak.adapters.spi.AuthChallenge;
import org.keycloak.adapters.spi.AuthOutcome;
import org.keycloak.adapters.spi.HttpFacade;
import org.keycloak.adapters.spi.UserSessionManagement;
import org.keycloak.common.VerificationException;
import org.keycloak.common.util.Base64Url;
import org.keycloak.common.util.SecretGenerator;
import org.keycloak.enums.TokenStore;
import org.keycloak.jose.jws.JWSInput;
import org.keycloak.jose.jws.JWSInputException;
import org.keycloak.representations.AccessTokenResponse;

import javax.servlet.Filter;
import javax.servlet.FilterChain;
import javax.servlet.ServletException;
import javax.servlet.ServletRequest;
import javax.servlet.ServletResponse;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpSession;
import java.io.IOException;
import java.net.MalformedURLException;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.util.List;
import java.util.Map;
import java.util.logging.Level;
import java.util.logging.Logger;

public class CustomKeycloakOIDCFilterWithPkceSupport extends KeycloakOIDCFilter {

    private final static Logger log = Logger.getLogger(CustomKeycloakOIDCFilterWithPkceSupport.class.getName());
    public static final String PKCE_CODE_VERIFIER_SESSION_ATTRIBUTE = "pkceCodeVerifier";
    public static final String PKCE_CODE_CHALLENGE_SESSION_ATTRIBUTE = "pkceCodeChallenge";

    @Override
    public void doFilter(ServletRequest req, ServletResponse res, FilterChain chain) throws IOException, ServletException {
        log.fine("Keycloak OIDC Filter");
        HttpServletRequest request = (HttpServletRequest) req;
        HttpServletResponse response = (HttpServletResponse) res;

        if (shouldSkip(request)) {
            chain.doFilter(req, res);
            return;
        }

        OIDCServletHttpFacade facade = new OIDCServletHttpFacade(request, response);
        KeycloakDeployment deployment = deploymentContext.resolveDeployment(facade);
        if (deployment == null || !deployment.isConfigured()) {
            response.sendError(403);
            log.fine("deployment not configured");
            return;
        }

        PreAuthActionsHandler preActions = new PreAuthActionsHandler(new IdMapperUserSessionManagement(), deploymentContext, facade);

        if (preActions.handleRequest()) {
            //System.err.println("**************** preActions.handleRequest happened!");
            return;
        }


        nodesRegistrationManagement.tryRegister(deployment);
        OIDCFilterSessionStore tokenStore = new OIDCFilterSessionStore(request, facade, 100000, deployment, idMapper);
        tokenStore.checkCurrentToken();


        FilterRequestAuthenticator authenticator = new CustomFilterRequestAuthenticatorWithPkceSupport(deployment, tokenStore, facade, request, 8443);
        AuthOutcome outcome = authenticator.authenticate();
        if (outcome == AuthOutcome.AUTHENTICATED) {
            log.fine("AUTHENTICATED");
            if (facade.isEnded()) {
                return;
            }
            AuthenticatedActionsHandler actions = new AuthenticatedActionsHandler(deployment, facade);
            if (actions.handledRequest()) {
                return;
            } else {
                HttpServletRequestWrapper wrapper = tokenStore.buildWrapper();
                chain.doFilter(wrapper, res);
                return;
            }
        }
        AuthChallenge challenge = authenticator.getChallenge();
        if (challenge != null) {
            log.fine("challenge");
            challenge.challenge(facade);
            return;
        }
        response.sendError(403);

    }



    /**
     * Decides whether this {@link Filter} should skip the given {@link HttpServletRequest} based on the configured {@link KeycloakOIDCFilter#skipPattern}.
     * Patterns are matched against the {@link HttpServletRequest#getRequestURI() requestURI} of a request without the context-path.
     * A request for {@code /myapp/index.html} would be tested with {@code /index.html} against the skip pattern.
     * Skipped requests will not be processed further by {@link KeycloakOIDCFilter} and immediately delegated to the {@link FilterChain}.
     *
     * @param request the request to check
     * @return {@code true} if the request should not be handled,
     * {@code false} otherwise.
     */
    private boolean shouldSkip(HttpServletRequest request) {

        if (skipPattern == null) {
            return false;
        }

        String requestPath = request.getRequestURI().substring(request.getContextPath().length());
        return skipPattern.matcher(requestPath).matches();
    }

    public static class Pkce {
        // https://tools.ietf.org/html/rfc7636#section-4.1
        public static final int PKCE_CODE_VERIFIER_MAX_LENGTH = 128;

        private final String codeChallenge;
        private final String codeVerifier;

        public Pkce(String codeVerifier, String codeChallenge) {
            this.codeChallenge = codeChallenge;
            this.codeVerifier = codeVerifier;
        }

        public String getCodeChallenge() {
            return codeChallenge;
        }

        public String getCodeVerifier() {
            return codeVerifier;
        }

        public static Pkce generatePkce() {
            try {
                String codeVerifier = SecretGenerator.getInstance().randomString(PKCE_CODE_VERIFIER_MAX_LENGTH);
                String codeChallenge = generateS256CodeChallenge(codeVerifier);
                return new Pkce(codeVerifier, codeChallenge);
            } catch (Exception ex){
                throw new RuntimeException("Could not generate PKCE", ex);
            }
        }

        // https://tools.ietf.org/html/rfc7636#section-4.6
        private static String generateS256CodeChallenge(String codeVerifier) throws Exception {
            MessageDigest md = MessageDigest.getInstance("SHA-256");
            md.update(codeVerifier.getBytes(StandardCharsets.ISO_8859_1));
            return Base64Url.encode(md.digest());
        }
    }

    private class IdMapperUserSessionManagement implements UserSessionManagement {
        @Override
        public void logoutAll() {
            if (idMapper != null) {
                idMapper.clear();
            }
        }

        @Override
        public void logoutHttpSessions(List<String> ids) {
            log.fine("**************** logoutHttpSessions");
            //System.err.println("**************** logoutHttpSessions");
            for (String id : ids) {
                log.finest("removed idMapper: " + id);
                idMapper.removeSession(id);
            }

        }
    }

    static class CustomFilterRequestAuthenticatorWithPkceSupport extends FilterRequestAuthenticator {

        public CustomFilterRequestAuthenticatorWithPkceSupport(KeycloakDeployment deployment, AdapterTokenStore tokenStore, OIDCHttpFacade facade, HttpServletRequest request, int sslRedirectPort) {
            super(deployment, tokenStore, facade, request, sslRedirectPort);
        }

        @Override
        protected OAuthRequestAuthenticator createOAuthAuthenticator() {
            if (deployment.isPkce()) {

                HttpSession session = request.getSession();
                if (session.getAttribute(PKCE_CODE_CHALLENGE_SESSION_ATTRIBUTE) == null) {
                    // no PKCE challenge present yet in session, so we need to generate one
                    Pkce pkce = Pkce.generatePkce();
                    session.setAttribute(PKCE_CODE_VERIFIER_SESSION_ATTRIBUTE, pkce.getCodeVerifier());
                    session.setAttribute(PKCE_CODE_CHALLENGE_SESSION_ATTRIBUTE, pkce.getCodeChallenge());
                }

                return new CustomOAuthRequestAuthenticator(this, facade, deployment, sslRedirectPort, tokenStore, request.getSession());
            }
            return super.createOAuthAuthenticator();
        }

        public String changeHttpSessionId(boolean create) {
            return super.changeHttpSessionId(create);
        }
    }

    static class CustomOAuthRequestAuthenticator extends OAuthRequestAuthenticator {

        private static final Logger log = Logger.getLogger(OAuthRequestAuthenticator.class.getName());

        private final HttpSession session;

        public CustomOAuthRequestAuthenticator(RequestAuthenticator requestAuthenticator, HttpFacade facade, KeycloakDeployment deployment, int sslRedirectPort, AdapterSessionStore tokenStore, HttpSession session) {
            super(requestAuthenticator, facade, deployment, sslRedirectPort, tokenStore);
            this.session = session;
        }

        @Override
        protected String getRedirectUri(String state) {
            String redirectUri = super.getRedirectUri(state);

            // add PKCE
            // add code_challenge
            // add code_challenge_method

            String pkceCodeChallenge = String.valueOf(session.getAttribute(PKCE_CODE_CHALLENGE_SESSION_ATTRIBUTE));
            try {
                redirectUri += "&code_challenge=" + pkceCodeChallenge + "&code_challenge_method=S256" ;
            } catch (Exception ignore) {
                // throw new RuntimeException(e);
            }

            return redirectUri;
        }

        /**
         * Start or continue the oauth login process.
         * <p/>
         * if code query parameter is not present, then browser is redirected to authUrl.  The redirect URL will be
         * the URL of the current request.
         * <p/>
         * If code query parameter is present, then an access token is obtained by invoking a secure request to the codeUrl.
         * If the access token is obtained, the browser is again redirected to the current request URL, but any OAuth
         * protocol specific query parameters are removed.
         *
         * @return null if an access token was obtained, otherwise a challenge is returned
         */
        protected AuthChallenge resolveCode(String code) {
            // abort if not HTTPS
            if (!isRequestSecure() && deployment.getSslRequired().isRequired(facade.getRequest().getRemoteAddr())) {
                log.log(Level.SEVERE,"Adapter requires SSL. Request: " + facade.getRequest().getURI());
                return challenge(403, OIDCAuthenticationError.Reason.SSL_REQUIRED, null);
            }

            log.log(Level.FINEST,"checking state cookie for after code");
            AuthChallenge challenge = checkStateCookie();
            if (challenge != null) return challenge;

            AccessTokenResponse tokenResponse = null;
            strippedOauthParametersRequestUri = rewrittenRedirectUri(stripOauthParametersFromRedirect());

            try {
                // For COOKIE store we don't have httpSessionId and single sign-out won't be available
                String httpSessionId = deployment.getTokenStore() == TokenStore.SESSION ? ((CustomFilterRequestAuthenticatorWithPkceSupport)reqAuthenticator).changeHttpSessionId(true) : null;
                if (deployment.isPkce()) {
                    String pkceCodeVerifier = String.valueOf(session.getAttribute(PKCE_CODE_VERIFIER_SESSION_ATTRIBUTE));
                    tokenResponse = ServerRequest.invokeAccessCodeToToken(deployment, code, strippedOauthParametersRequestUri, httpSessionId, pkceCodeVerifier);
                } else {
                    tokenResponse = ServerRequest.invokeAccessCodeToToken(deployment, code, strippedOauthParametersRequestUri, httpSessionId);
                }
            } catch (ServerRequest.HttpFailure failure) {
                log.log(Level.SEVERE,"failed to turn code into token");
                log.log(Level.SEVERE, "status from server: " + failure.getStatus());
                if (failure.getError() != null && !failure.getError().trim().isEmpty()) {
                    log.log(Level.SEVERE, "   " + failure.getError());
                }
                return challenge(403, OIDCAuthenticationError.Reason.CODE_TO_TOKEN_FAILURE, null);

            } catch (IOException e) {
                log.log(Level.SEVERE, "failed to turn code into token", e);
                return challenge(403, OIDCAuthenticationError.Reason.CODE_TO_TOKEN_FAILURE, null);
            } finally {
                session.removeAttribute(PKCE_CODE_VERIFIER_SESSION_ATTRIBUTE);
                session.removeAttribute(PKCE_CODE_CHALLENGE_SESSION_ATTRIBUTE);
            }

            tokenString = tokenResponse.getToken();
            refreshToken = tokenResponse.getRefreshToken();
            idTokenString = tokenResponse.getIdToken();

            log.log(Level.FINEST, "Verifying tokens");
            if (log.isLoggable(Level.FINEST)) {
                logToken("\taccess_token", tokenString);
                logToken("\tid_token", idTokenString);
                logToken("\trefresh_token", refreshToken);
            }

            try {
                AdapterTokenVerifier.VerifiedTokens tokens = AdapterTokenVerifier.verifyTokens(tokenString, idTokenString, deployment);
                token = tokens.getAccessToken();
                idToken = tokens.getIdToken();
                log.log(Level.FINEST,"Token Verification succeeded!");
            } catch (VerificationException e) {
                log.log(Level.SEVERE,"failed verification of token: " + e.getMessage());
                return challenge(403, OIDCAuthenticationError.Reason.INVALID_TOKEN, null);
            }
            if (tokenResponse.getNotBeforePolicy() > deployment.getNotBefore()) {
                deployment.updateNotBefore(tokenResponse.getNotBeforePolicy());
            }
            if (token.getIssuedAt() < deployment.getNotBefore()) {
                log.log(Level.SEVERE,"Stale token");
                return challenge(403, OIDCAuthenticationError.Reason.STALE_TOKEN, null);
            }
            log.log(Level.FINEST,"successful authenticated");
            return null;
        }

        private String rewrittenRedirectUri(String originalUri) {
            Map<String, String> rewriteRules = deployment.getRedirectRewriteRules();
            if(rewriteRules != null && !rewriteRules.isEmpty()) {
                try {
                    URL url = new URL(originalUri);
                    Map.Entry<String, String> rule =  rewriteRules.entrySet().iterator().next();
                    StringBuilder redirectUriBuilder = new StringBuilder(url.getProtocol());
                    redirectUriBuilder.append("://"+ url.getAuthority());
                    redirectUriBuilder.append(url.getPath().replaceFirst(rule.getKey(), rule.getValue()));
                    return redirectUriBuilder.toString();
                } catch (MalformedURLException ex) {
                    log.log(Level.SEVERE,"Not a valid request url");
                    throw new RuntimeException(ex);
                }
            }
            return originalUri;
        }

        private void logToken(String name, String token) {
            try {
                JWSInput jwsInput = new JWSInput(token);
                String wireString = jwsInput.getWireString();
                log.log(Level.FINEST, String.format("\t%s: %s", name, wireString.substring(0, wireString.lastIndexOf(".")) + ".signature"));
            } catch (JWSInputException e) {
                log.log(Level.SEVERE, String.format("Failed to parse %s: %s", name, token), e);
            }
        }
    }
}