keycloak-d4science-spi-parent/protocol-mapper/src/main/java/org/gcube/keycloak/protocol/oidc/mapper/D4ScienceContextMapper.java

package org.gcube.keycloak.protocol.oidc.mapper;

import java.util.ArrayList;
import java.util.List;

import org.jboss.logging.Logger;
import org.keycloak.models.ClientSessionContext;
import org.keycloak.models.KeycloakSession;
import org.keycloak.models.ProtocolMapperModel;
import org.keycloak.models.UserSessionModel;
import org.keycloak.protocol.oidc.mappers.AbstractOIDCProtocolMapper;
import org.keycloak.protocol.oidc.mappers.OIDCAccessTokenMapper;
import org.keycloak.protocol.oidc.mappers.OIDCAttributeMapperHelper;
import org.keycloak.provider.ProviderConfigProperty;
import org.keycloak.representations.AccessToken;
import org.keycloak.representations.AccessToken.Access;
import org.keycloak.representations.IDToken;

public class D4ScienceContextMapper extends AbstractOIDCProtocolMapper implements OIDCAccessTokenMapper {

    private static final Logger logger = Logger.getLogger(D4ScienceContextMapper.class);

    public static final String HTTP_REQUEST_HEADER_NAME = "d4scm.header-name";
    public static final String NARROW_RESOURCE_ACCESS = "d4scm.narrow-ra";

    private static final List<ProviderConfigProperty> CONFIG_PROPERTIES = new ArrayList<>();

    // Assuring that the mapper is executed as last
    private static final int PRIORITY = Integer.MAX_VALUE;
    private static final String DISPLAY_TYPE = "OIDC D4Science Context Mapper";
    private static final String PROVIDER_ID = "oidc-d4scince-context-mapper";

    public static final String DEFAULT_HEADER_NAME = "X-D4Science-Context";
    public static final String DEFAULT_TOKEN_CLAIM = "aud";

    static {
        OIDCAttributeMapperHelper.addTokenClaimNameConfig(CONFIG_PROPERTIES);
        CONFIG_PROPERTIES.forEach(conf -> {
            if (OIDCAttributeMapperHelper.TOKEN_CLAIM_NAME.equals(conf.getName()))
                conf.setDefaultValue(DEFAULT_TOKEN_CLAIM);
            conf.setReadOnly(true);
        });

        OIDCAttributeMapperHelper.addIncludeInTokensConfig(CONFIG_PROPERTIES, D4ScienceContextMapper.class);

        ProviderConfigProperty headerProperty = new ProviderConfigProperty();
        headerProperty.setName(HTTP_REQUEST_HEADER_NAME);
        headerProperty.setLabel("HTTP request header name with the requested context");
        headerProperty.setType(ProviderConfigProperty.STRING_TYPE);
        headerProperty.setHelpText("The HTTP header that contains the requested context to be mapped as the requested in the configured claim");
        headerProperty.setDefaultValue(DEFAULT_HEADER_NAME);
        headerProperty.setReadOnly(true);
        CONFIG_PROPERTIES.add(headerProperty);

        ProviderConfigProperty narrowProperty = new ProviderConfigProperty();
        narrowProperty.setName(NARROW_RESOURCE_ACCESS);
        narrowProperty.setLabel("Narrow down resource access array?");
        narrowProperty.setType(ProviderConfigProperty.BOOLEAN_TYPE);
        narrowProperty.setHelpText("Narrow down resource access claim to contain only the requested context entry");
        CONFIG_PROPERTIES.add(narrowProperty);
    }

    @Override
    public String getDisplayCategory() {
        return TOKEN_MAPPER_CATEGORY;
    }

    @Override
    public int getPriority() {
        return PRIORITY;
    }

    @Override
    public String getDisplayType() {
        return DISPLAY_TYPE;
    }

    @Override
    public String getHelpText() {
        return "Maps the D4Science context audience by reading the configured header's value and sets it as the configured token claim, if it is in scope";
    }

    @Override
    public List<ProviderConfigProperty> getConfigProperties() {
        return CONFIG_PROPERTIES;
    }

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

    @Override
    protected void setClaim(final IDToken token,
            final ProtocolMapperModel mappingModel,
            final UserSessionModel userSession,
            final KeycloakSession keycloakSession,
            final ClientSessionContext clientSessionCtx) {

        // Since only the OIDCAccessTokenMapper interface is implemented, we are almost sure that
        // the token object is an AccessToken but adding a specific check anyway
        if (token instanceof AccessToken) {
            AccessToken accessToken = ((AccessToken) token);
            String headerName = mappingModel.getConfig().get(HTTP_REQUEST_HEADER_NAME);
            if (headerName == null || "".equals(headerName)) {
                headerName = DEFAULT_HEADER_NAME;
            }
            logger.debugf("Looking for the '%s' header", headerName);
            String requestedD4SContext = keycloakSession.getContext().getRequestHeaders().getHeaderString(headerName);

            if (requestedD4SContext != null && !"".equals(requestedD4SContext)) {
                logger.debugf("Checking resource access for the requested context: %s", requestedD4SContext);
                Access contextAccessInResourceAccess = accessToken.getResourceAccess(requestedD4SContext);
                if (contextAccessInResourceAccess != null) {
                    logger.debugf("Mapping it as the configured claim: %s",
                            mappingModel.getConfig().get(OIDCAttributeMapperHelper.TOKEN_CLAIM_NAME));

                    OIDCAttributeMapperHelper.mapClaim(token, mappingModel, requestedD4SContext);

                    if (Boolean.parseBoolean(mappingModel.getConfig().get(NARROW_RESOURCE_ACCESS))) {
                        // Removing all access details but the requested context
                        accessToken.getResourceAccess().clear();
                        accessToken.getResourceAccess().put(requestedD4SContext, contextAccessInResourceAccess);
                    }
                } else {
                    logger.warnf("Requested context '%s' is not accessible to the client: %s", requestedD4SContext,
                            clientSessionCtx.getClientSession().getClient().getName());
                }
            }
        }
    }

}