Create custom JWT Token assertion in KeyCloak

We all know Keycloak can easily generate a JWT token and include specific claim values. These values can be either hardcoded or taken from the identity account.

For those who are not aware, a JWT token claim is information asserted about a subject and used to convey certain information between parties. Claims are key-value pairs that provide metadata about the token and its subject. They are part of the payload of a JWT token and can be used for authentication, authorization, and information exchange.

{
  "iss": "https://example.com",
  "sub": "1234567890",
  "aud": "https://example.com/api",
  "exp": 1618884471,
  "iat": 1618880871,
  "role": "admin",
  "email": "user@example.com",
  "department": "engineering"
}

My team manages multiple microservices, and there was an interesting request to include certain information in the JWT token and share it among all the microservices as if it were a stateless cache. This improves performance since those services won’t have to retrieve this information on their own. The sequence diagram would look like this:

While Keycloak does not support this functionality out of the box, its Service Provider Interface (SPI) allows us to write hooks using Java to extend and customize Keycloak’s functionality to achieve this.

In this blog, I want to show you how I implemented the above request by going through the following steps:

1. Create a New Maven Project:

   • Set up a new Maven project for our custom SPI mapper.

2. Implement the Custom Protocol Mapper:

   • Implement the interfaces to create our custom protocol mapper. Its responsibility will be to parse the incoming request parameter calls custom_claim into a set of key-pair values and add them into the JWT token as claims.

3. Deploy the Protocol Mapper:

   • Deploy our custom protocol mapper to the Keycloak server.

4. Register & Configure the Protocol Mapper:

   • Register our custom protocol mapper with Keycloak and configure KeyCloak to utilize our new custom protocol mapper.

Before we begin, I assume you already have a local KeyCloak server running and you have the latest Java SDK and Apache Maven installed.

Create a New Maven Project:

To create a Maven project we can utilize the following command in the command prompt window:

mvn archetype:generate -DgroupId=com.example.keycloak -DartifactId=custom-claim-spi -DarchetypeArtifactId=maven-archetype-quickstart -DinteractiveMode=false

This should create a project with some boiler plate with some file structure such as one below, to keep thing simple, we will remove the default App.java file and the entire test folder.

Implement the Custom Protocol Mapper:

• Create a file calls CustomClaimMapper.java under src\main\java\com\example directory.

• Create a file calls org.keycloak.protocol.ProtocolMapper under the folder structure resources\META-INF\services and in that file put com.example.keycloak.CustomClaimMapper

In this case, the file org.keycloak.protocol.ProtocolMapper contains the fully qualified name of a class (com.example.keycloak.CustomClaimMapper). This indicates that com.example.keycloak.CustomClaimMapper is an implementation of the ProtocolMapper service interface. When the service is requested, the Java runtime will load this implementation. The new project structure will now looks like:

CustomClaimMapper.java is our main class where we will implement this custom mapper. It needs to implement the AbstractOIDCProtocolMapper class. it is an abstract class provided by Keycloak that simplifies the implementation of custom protocol mappers for the OpenID Connect (OIDC) protocol. It provides a base implementation for common functionality required by OIDC protocol mappers.

Now, the function we want to override and where most of our logic will be setClaim . The setClaim method provides several important inputs that we can utilize to customize the token. Here are the input parameters to this function :

• IDToken token: This is the token being generated (e.g., ID token, access token, or refresh token). You can modify this token by adding custom claims to it.

• ProtocolMapperModel mappingModel: This model contains the configuration of the protocol mapper. It includes any settings or parameters defined for the mapper in the Keycloak admin console.

• UserSessionModel userSession: This represents the user's session. It provides access to user-related information, such as user attributes, roles, and session details.

• KeycloakSession keycloakSession: This is the overall KeyCloak session. It provides access to various Keycloak services and contexts, such as the HTTP request context, which can be used to extract additional information from the request.

• ClientSessionContext clientSessionCtx: This represents the client session context. It provides information about the client application requesting the token, including client attributes and roles.

With this in mind our function implementation would be:

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

            // extract the request parameters from the token request which is part of the
            // keycloak session object.
            KeycloakContext keycloakContext = keycloakSession.getContext();
            keycloakContext.getHttpRequest().getMultiPartFormParameters().forEach((k, v) -> {
                logger.debug("k:" + k + ", v:" + v);
            });

            keycloakContext.getHttpRequest()
                    .getMultiPartFormParameters()
                    .get(CUSTOM_CLAIM_PARAM_KEY);

            // Extract the multipart form parameters
            MultivaluedMap<String, String> formParameters = keycloakContext.getHttpRequest().getDecodedFormParameters();

            // Extract the value of the request parameter.
            List<String> customClaimValues = formParameters.get(CUSTOM_CLAIM_PARAM_KEY);

            if (customClaimValues != null && !customClaimValues.isEmpty()) {
                // Assuming FormatPartValue has a method getValue() that returns the string
                // value
                String customClaimValue = customClaimValues.get(0);
                logger.debug("request claim: " + customClaimValue);

                // Parse the JSON value of the parameter key
                ObjectMapper objectMapper = new ObjectMapper();
                JsonNode jsonNode = objectMapper.readTree(customClaimValue);

                // Iterate over the JSON object and add each key-value pair
                // as a claim key and claim value to the JWT Token.
                Iterator<Map.Entry<String, JsonNode>> fields = jsonNode.fields();
                while (fields.hasNext()) {
                    Map.Entry<String, JsonNode> field = fields.next();
                    String key = field.getKey();
                    String value = field.getValue().asText();
                    logger.info("Claim key: " + key + ", Claim value: " + value);

                    // Add the extracted key-value pairs to the token
                    token.getOtherClaims().put(key, value);
                }
            }
        } catch (Exception e) {
            logger.error("Error setting custom claim", e);
        }
    }

In the provided code snippet, the setClaim method performs the following steps:

1. Extracting Request Parameters: The method extracts request parameters from the HTTP request context using keycloakSession.getContext().getHttpRequest().getMultiPartFormParameters(). This allows you to access any custom parameters sent in the token request.

2. Extracting Form Parameters: The method extracts form parameters from the HTTP request using keycloakContext.getHttpRequest().getDecodedFormParameters(). This provides access to any form data submitted with the request.

3. Retrieving Custom Claim Values: The method retrieves the value of a specific custom claim parameter (CUSTOM_CLAIM_PARAM_KEY) from the form parameters. If the parameter is present, it extracts its value.

4. Adding Custom Claim to Token: If the custom claim value is found, the method adds it to the token. This is done by modifying the token object, which will include the custom claim in the final token issued by Keycloak.

Once the code is ready, you can run mvn compile to build the project and place the compiled .class files in the classes directory. To use this in Keycloak, we must first package it into a JAR file. Run the command mvn package to package the compiled classes and resources into a JAR file, which will be placed in the target directory.

Deploy the Protocol Mapper:

In order to deploy this custom protocol mapper to Keyloak we must take the JAR file generated from the prior step and place it into KeyCloak’s providers folder.

Next restart your KeyCloak server and you should see KeyCloak picks up this mapper during its startup.

Register & Configure the Protocol Mapper:

Now that Keycloak is able to load this custom mapper, we need to configure it to be used with a specific Keycloak client so that when we request a JWT token using that specific client ID, Keycloak can perform these steps to add in those custom claims.

1. Log in to the Keycloak admin console.

2. Click on "Client Scopes" and then click "Create client scope" and give it a meaningful name.

3. Once the scope is created, click on the "Mappers" tab and then click "Create" to configure a new mapper. In the pop-up window that opens, you should see the Custom Token Claim Mapper that we just deployed.

4. 

   Click Save

5. 

   Go back to Clients on the left menu and select the client you wish to apply this custom mapper to. On the client, select the Client scopes tab and click Add client scope and select the scope you just created and click Add.

   

   

   Once that is done, Keycloak is ready to work with the custom mapper we configured. To test this out, I will use Postman to send a token request to Keycloak using the Client Credentials Flow. In that request, I will include a parameter that contains a JSON object with key-value pairs. What we should get back is the standard JWT token we expect, but that token should also contain the key-value pairs we included in our request.

If you copy the generated JWT Token and check it at the site https://jwt.io/ , you should see in the decoded portion of the token’s payload, the values we sent in as the custom_claim parameter

This enhancement allows Keycloak to dynamically parse any request parameters and include them in the generated token to be issued. Without this custom mapper, Keycloak can only include existing identity metadata it has. The current method allows Keycloak to include any data provided in the request, without needing prior knowledge of that data.

I hope this shows you the flexibility which KeyCloak SPI brings and the freedom you can play around with the JWT Token payload. As always the code and Postman collection for this can be found on my GitHub Repo.