Custom CustomAzureOidcIdentityProvider to support client assertions with managed identities

CustomAzureOidcIdentityProvider.java

package com.github.thomasdarimont.keycloak.custom.idp.azure;

import com.google.auto.service.AutoService;
import lombok.val;
import org.keycloak.authentication.ClientAuthenticationFlowContext;
import org.keycloak.authentication.authenticators.client.FederatedJWTClientValidator;
import org.keycloak.broker.oidc.OIDCIdentityProvider;
import org.keycloak.broker.oidc.OIDCIdentityProviderConfig;
import org.keycloak.broker.oidc.OIDCIdentityProviderFactory;
import org.keycloak.broker.provider.IdentityProviderFactory;
import org.keycloak.common.Profile;
import org.keycloak.models.IdentityProviderModel;
import org.keycloak.models.KeycloakSession;

import java.time.Duration;
import java.util.ArrayList;
import java.util.List;

public class CustomAzureOidcIdentityProvider extends OIDCIdentityProvider {

    public static final String PROVIDER_ID = "acme-oidc";

    public CustomAzureOidcIdentityProvider(KeycloakSession session, OIDCIdentityProviderConfig config) {
        super(session, config);
    }

    @Override
    public boolean verifyClientAssertion(ClientAuthenticationFlowContext context) throws Exception {
        OIDCIdentityProviderConfig config = getConfig();

        FederatedJWTClientValidator validator = new CustomFederatedJWTClientValidator(context, v -> verifySignature(v.getJws()),
                config.getIssuer(), config.getAllowedClockSkew(), config.isSupportsClientAssertionReuse());
        // Azure Managed Identity Access Token are valid for 55min, default in Keycloak is 5mins
        validator.setMaximumExpirationTime((int)Duration.ofMinutes(120).toSeconds());

        if (!Profile.isFeatureEnabled(Profile.Feature.CLIENT_AUTH_FEDERATED)) {
            return false;
        }

        if (!config.isSupportsClientAssertions()) {
            throw new RuntimeException("Issuer does not support client assertions");
        }

        return validator.validate();
    }

    public static class CustomFederatedJWTClientValidator extends FederatedJWTClientValidator {

        public CustomFederatedJWTClientValidator(ClientAuthenticationFlowContext context, SignatureValidator signatureValidator, String expectedTokenIssuer, int allowedClockSkew, boolean reusePermitted) throws Exception {
            super(context, signatureValidator, expectedTokenIssuer, allowedClockSkew, reusePermitted);
        }

        @Override
        protected List<String> getExpectedAudiences() {
            List<String> audiences = new ArrayList<>(super.getExpectedAudiences());
            // TODO make additional allowed audiences configurable
            audiences.add("api://someuuid");
            audiences.add("https://somesubdomain.sometenanttrusted.custom.domain.com");
            return audiences;
        }
    }

    @AutoService(IdentityProviderFactory.class)
    public static class Factory extends OIDCIdentityProviderFactory {
        @Override
        public String getId() {
            return CustomAzureOidcIdentityProvider.PROVIDER_ID;
        }

        @Override
        public String getName() {
            return "Acme: Azure OpenID Connect v1.0";
        }

        @Override
        public OIDCIdentityProvider create(KeycloakSession session, IdentityProviderModel model) {
            return new CustomAzureOidcIdentityProvider(session, new OIDCIdentityProviderConfig(model));
        }
    }
}