DynamicIdpRedirectAuthenticator with configurable fallback

DynamicIdpRedirectAuthenticator.java

package com.github.thomasdarimont.keycloak.auth.dynamicidp;

import lombok.extern.jbosslog.JBossLog;
import org.keycloak.OAuth2Constants;
import org.keycloak.authentication.AuthenticationFlowContext;
import org.keycloak.authentication.AuthenticationFlowError;
import org.keycloak.authentication.Authenticator;
import org.keycloak.constants.AdapterConstants;
import org.keycloak.events.Errors;
import org.keycloak.models.AuthenticatorConfigModel;
import org.keycloak.models.IdentityProviderModel;
import org.keycloak.models.KeycloakSession;
import org.keycloak.models.RealmModel;
import org.keycloak.models.UserModel;
import org.keycloak.services.Urls;
import org.keycloak.services.managers.ClientSessionCode;

import javax.ws.rs.core.Response;
import javax.ws.rs.core.UriBuilder;
import java.net.URI;
import java.util.List;
import java.util.Map;
import java.util.function.Function;

@JBossLog
public class DynamicIdpRedirectAuthenticator implements Authenticator {

    public static final String TARGET_IDP_ATTRIBUTE = "targetIdp";

    static final String EMAIL_TO_IDP_MAPPING_CONFIG_PROPERTY = "email-to-idp-mapping";

    static final String FALLBACK_TO_AUTHFLOW_CONFIG_PROPERTY = "fallback-to-authflow";

    private final KeycloakSession session;

    public DynamicIdpRedirectAuthenticator(KeycloakSession session) {
        this.session = session;
    }

    @Override
    public void authenticate(AuthenticationFlowContext context) {

        UserModel user = context.getUser();
        if (user == null) {
            context.attempted();
            return;
        }

        String targetIdp = determineTargetIdp(user, context);
        if (targetIdp != null) {
            redirect(context, targetIdp);
            return;
        }

        boolean fallbackToAuthFlow = getConfigValueOrDefault(context.getAuthenticatorConfig(), FALLBACK_TO_AUTHFLOW_CONFIG_PROPERTY, "true", Boolean::parseBoolean);
        if (fallbackToAuthFlow) {
            context.attempted();
            return;
        }

        context.getEvent().error(Errors.UNKNOWN_IDENTITY_PROVIDER);
        context.failure(AuthenticationFlowError.IDENTITY_PROVIDER_NOT_FOUND);
        context.cancelLogin();
        context.resetFlow();
    }

    private void redirect(AuthenticationFlowContext context, String providerId) {

        IdentityProviderModel identityProviderModel = selectIdp(context, providerId);
        if (identityProviderModel == null || !identityProviderModel.isEnabled()) {
            log.warnf("Provider not found or not enabled for realm %s", providerId);
            context.attempted();
            return;
        }

        String accessCode = new ClientSessionCode<>(context.getSession(), context.getRealm(), context.getAuthenticationSession()).getOrGenerateCode();
        String clientId = context.getAuthenticationSession().getClient().getClientId();
        String tabId = context.getAuthenticationSession().getTabId();
        URI location = Urls.identityProviderAuthnRequest(context.getUriInfo().getBaseUri(), providerId, context.getRealm().getName(), accessCode, clientId, tabId);
        if (context.getAuthenticationSession().getClientNote(OAuth2Constants.DISPLAY) != null) {
            location = UriBuilder.fromUri(location).queryParam(OAuth2Constants.DISPLAY, context.getAuthenticationSession().getClientNote(OAuth2Constants.DISPLAY)).build();
        }
        log.debugf("Redirecting to %s", providerId);
        Response response = Response.seeOther(location).build();
        context.forceChallenge(response);
    }

    private IdentityProviderModel selectIdp(AuthenticationFlowContext context, String providerId) {

        List<IdentityProviderModel> identityProviders = context.getRealm().getIdentityProviders();
        for (IdentityProviderModel identityProvider : identityProviders) {

            if (!identityProvider.isEnabled()) {
                continue;
            }

            if (providerId.equals(identityProvider.getAlias())) {
                return identityProvider;
            }
        }

        return null;
    }


    private String determineTargetIdp(UserModel user, AuthenticationFlowContext context) {

        String targetIdp = determineTargetIdpFromUrlParameter(context);
        if (targetIdp != null) {
            return targetIdp;
        }

        targetIdp = determineTargetIdpViaAttribute(user);
        if (targetIdp != null) {
            return targetIdp;
        }

        return determineTargetIdpViaUserEmail(user, context);
    }

    private String determineTargetIdpFromUrlParameter(AuthenticationFlowContext context) {
        return context.getUriInfo().getQueryParameters().getFirst(AdapterConstants.KC_IDP_HINT);
    }

    private String determineTargetIdpViaAttribute(UserModel user) {
        return user.getFirstAttribute(TARGET_IDP_ATTRIBUTE);
    }

    private String determineTargetIdpViaUserEmail(UserModel user, AuthenticationFlowContext context) {

        String email = user.getEmail();
        if (email == null) {
            return null;
        }

        String mappingString = getConfigValueOrDefault(context.getAuthenticatorConfig(), EMAIL_TO_IDP_MAPPING_CONFIG_PROPERTY, "", String::valueOf);
        String[] mappings = mappingString.split(";");
        for (String mapping : mappings) {
            String[] emailSuffixPatternToIdpId = mapping.split("/");
            String emailSuffixPattern = emailSuffixPatternToIdpId[0];
            String idpId = emailSuffixPatternToIdpId[1];

            if (email.matches(emailSuffixPattern)) {
                return idpId;
            }
        }

        return null;
    }

    private <T> T getConfigValueOrDefault(AuthenticatorConfigModel configModel, String key, String defaultValue, Function<String, T> converter) {

        if (configModel == null) {
            return converter.apply(defaultValue);
        }

        Map<String, String> config = configModel.getConfig();
        if (config == null || config.isEmpty()) {
            return converter.apply(defaultValue);
        }

        return converter.apply(config.getOrDefault(key, defaultValue));
    }

    @Override
    public void action(AuthenticationFlowContext context) {
        // NOOP
    }

    @Override
    public boolean requiresUser() {
        return true;
    }

    @Override
    public boolean configuredFor(KeycloakSession session, RealmModel realm, UserModel user) {
        return true;
    }

    @Override
    public void setRequiredActions(KeycloakSession session, RealmModel realm, UserModel user) {
        // NOOP
    }

    @Override
    public void close() {
        // NOOP
    }
}

DynamicIdpRedirectAuthenticatorFactory.java

package com.github.thomasdarimont.keycloak.auth.dynamicidp;

import com.google.auto.service.AutoService;
import org.keycloak.Config;
import org.keycloak.authentication.Authenticator;
import org.keycloak.authentication.AuthenticatorFactory;
import org.keycloak.models.AuthenticationExecutionModel;
import org.keycloak.models.KeycloakSession;
import org.keycloak.models.KeycloakSessionFactory;
import org.keycloak.provider.ProviderConfigProperty;

import java.util.Arrays;
import java.util.List;

@AutoService(AuthenticatorFactory.class)
public class DynamicIdpRedirectAuthenticatorFactory implements AuthenticatorFactory {

    private static final String PROVIDER_ID = "auth-dynamic-idp-redirector";

    @Override
    public String getDisplayType() {
        return "Dynamic IDP Redirector";
    }

    @Override
    public String getReferenceCategory() {
        return null;
    }

    @Override
    public boolean isConfigurable() {
        return true;
    }

    public static final AuthenticationExecutionModel.Requirement[] REQUIREMENT_CHOICES = {
            AuthenticationExecutionModel.Requirement.REQUIRED, AuthenticationExecutionModel.Requirement.DISABLED
    };

    @Override
    public AuthenticationExecutionModel.Requirement[] getRequirementChoices() {
        return REQUIREMENT_CHOICES;
    }

    @Override
    public boolean isUserSetupAllowed() {
        return false;
    }

    @Override
    public String getHelpText() {
        return "Dynamic IDP Redirector";
    }

    @Override
    public List<ProviderConfigProperty> getConfigProperties() {

        ProviderConfigProperty emailToIdpMapping = new ProviderConfigProperty();
        emailToIdpMapping.setType(ProviderConfigProperty.STRING_TYPE);
        emailToIdpMapping.setName(DynamicIdpRedirectAuthenticator.EMAIL_TO_IDP_MAPPING_CONFIG_PROPERTY);
        emailToIdpMapping.setLabel("Email IDP Mapping");
        emailToIdpMapping.setHelpText("Email Suffix pattern to IDP Mapping. email-suffix/idp-id, multiple patterns can be delimited via ';', c.f.: example.com/idp1;.*foo.com/idp2;.*bar.(com|de)/idp3");

        ProviderConfigProperty fallbackToAuthFlow = new ProviderConfigProperty();
        fallbackToAuthFlow.setType(ProviderConfigProperty.BOOLEAN_TYPE);
        fallbackToAuthFlow.setName(DynamicIdpRedirectAuthenticator.FALLBACK_TO_AUTHFLOW_CONFIG_PROPERTY);
        fallbackToAuthFlow.setLabel("Fallback to Authflow");
        fallbackToAuthFlow.setHelpText("Fall back to Authflow if no target IdP could be identified.");
        fallbackToAuthFlow.setDefaultValue(true);

        return Arrays.asList(emailToIdpMapping, fallbackToAuthFlow);
    }

    @Override
    public void close() {
        // NOOP
    }

    @Override
    public Authenticator create(KeycloakSession session) {
        return new DynamicIdpRedirectAuthenticator(session);
    }

    @Override
    public void init(Config.Scope config) {
        // NOOP
    }

    @Override
    public void postInit(KeycloakSessionFactory factory) {
        // NOOP
    }

    @Override
    public String getId() {
        return PROVIDER_ID;
    }
}