• oauth

How to add custom properties to JWT tokens in Spring Security.

(Jama Software is a fast-growing company, in Portland, OR, building software for better, faster requirements definition, management, verification and validation.) In being a SaaS company, we are gradually chipping away at our good old monolith, turning pieces into micro-services that can scale horizontally, and that scale efficiently by use of multi-tenancy. A single micro-service can have multiple instances, and each instance serves a multitude of customers. Multi-tenancy has implications on application state, and a common pattern is for database tables to be shared across tenants, where each record links to a specific tenant. It also requires some lifting to make sure that the application always understands which tenant it is working for.

We are a Spring shop, and happy users of Spring Boot for our micro-services. We recently built the “Jama OAuth service”, which is an OAuth 2 compatible authorization server, that essentially issues access tokens to clients of our system (given their credentials). It implements OAuth’s so-called “client credentials” flow/grant type.

The access tokens are used to protect some REST resources. It was a must have requirement that the Jama OAuth service would support multi-tenancy. It was a natural choice to look at Spring Security, specifically Spring Security OAuth. This library however, does not, out of the box, support multi-tenancy.

When issuing access tokens, it is an interesting option to use JWT tokens. As this link shows, JWT tokens are not just a unique identifier by which the authorization server can verify your claim, rather they do include the details of the claim. In our case we could include not only information about the client, but also about what tenant they belong to. Our tokens could look something like below, which is awesome, because it would mean that resources in our ecosystem can validate an access token (extract all the information that they need) without having to contact the Jama OAuth service, an enormous performance gain.

{
    "exp": 12345,
    "scope": [
        "read"
    ],
    "tenant": "tenant1234",
    "jti": "d08e06d9-7408-4a01-bcf0-409ad23391ce",
    "client_id": "test1234"
}

This is almost exactly a JWT token that Spring Security OAuth spits out, using its JwtAccessTokenConverter, except for the custom property “tenant“. Unfortunately, there is no clear extension point to add custom properties. Conversely, when using Spring Security to validate an access token, what it gives your application code access to is an OAuth2Request, that does not include any custom properties. There is one more piece of tenant-awareness; in all we needed to address the following:

1. Make sure that the application understands which is the right tenant when an access token is requested. This one is simple for us: we have standardized on including an HTTP request header that identifies the tenant. If you forget to include this header, you are awarded an error message. If you include this header, we can use it together with your provided user name and password, to authenticate your request. We would then proceed to return you an access token (JWT token), see the first list item here.
2. Add custom property “tenant” to JWT tokens.
3. Read custom property “tenant” from JWT tokens and make it available to our application code.

Add Custom Property

Creating tokens is a function of the authentication server (in our case the “Jama OAuth service”). JWT tokens are generated in Spring by the JwtAccessTokenConverter. So, of course we override that class to get our way. It is being configured in our Spring JavaConfig as follows:

@Bean
public JwtAccessTokenConverter jwtAccessTokenConverter() throws Exception {
    // specifically the following line:
    JwtAccessTokenConverter converter = new TenantAwareJwtAccessTokenConverter();
    converter.setKeyPair(keyPair);
    return converter;
}

Our custom implementation starts as follows:

class TenantAwareJwtAccessTokenConverter extends JwtAccessTokenConverter { ...

Inside that class we retrieve the details of our client, which include the tenant, which we can then add to the access token:

@Override
public OAuth2AccessToken enhance(OAuth2AccessToken accessToken, OAuth2Authentication authentication) {
    ClientEntity clientEntity = getClientEntity(authentication);
    Map<String, Object> info = new LinkedHashMap<>(accessToken.getAdditionalInformation());
    info.putAll(clientEntity.getAdditionalInformationForToken()); // the additional information includes "tenant"="..."
    DefaultOAuth2AccessToken customAccessToken = new DefaultOAuth2AccessToken(accessToken);
    customAccessToken.setAdditionalInformation(info);
    return super.enhance(customAccessToken, authentication);
}

When retrieving the details of our client we take the client ID given by Spring, and combine it with the tenant header from the request (that we require users to include when offering their client credentials).

private ClientEntity getClientEntity(OAuth2Authentication authentication) {
    String clientId = (String) authentication.getPrincipal();
    String tenant = TenantHeaderHelper.getTenantFromRequest();
    return getClientEntityFromDatabase(clientId, tenant); // this includes some assertions to make sure the requested client exists
}

Read Custom Property

This assumes that your resource server is also using Spring Security. It may or may not be the same component as your authentication server (in our case the “Jama OAuth service”). JWT tokens are processed in Spring by a small army of classes, but we chose to override the DefaultAccessTokenConverter. This needs to be injected into the JwtAccessTokenConverter, here of course our own TenantAwareJwtAccessTokenConverter. It is being configured in our Spring JavaConfig as follows:

@Autowired
public void setJwtAccessTokenConverter(JwtAccessTokenConverter jwtAccessTokenConverter) {
    jwtAccessTokenConverter.setAccessTokenConverter(defaultAccessTokenConverter());
}

@Bean
DefaultAccessTokenConverter defaultAccessTokenConverter() {
    return new TenantAwareAccessTokenConverter();
}

Our custom implementation starts as follows:

class TenantAwareAccessTokenConverter extends DefaultAccessTokenConverter { ... 

Inside that class we can get access to a map that contains the raw data extracted from the JWT token, before Spring throws out our custom properties. Note that the super implementation already returns an OAuth2Authentication object. Inside that object, we substitute the originalOAuth2Request with our custom TenantAwareOAuth2Request.

@Override
public OAuth2Authentication extractAuthentication(Map<String, ?> map) {
    OAuth2Authentication authentication = super.extractAuthentication(map);
    TenantAwareOAuth2Request tenantAwareOAuth2Request = new TenantAwareOAuth2Request(authentication.getOAuth2Request());
    tenantAwareOAuth2Request.setTenant((String) map.get("tenant"));
    return new OAuth2Authentication(tenantAwareOAuth2Request, authentication.getUserAuthentication());
}

Our custom TenantAwareOAuth2Request looks as follows. Thanks to a useful constructor in the base class (“copy constructor”) our custom class remains relatively simple.

/**
 * Add a tenant to the existing {@link OAuth2Request}.
 */
public class extends OAuth2Request {
    public TenantAwareOAuth2Request(OAuth2Request other) {
        super(other);
    }

    private String tenant;

    public void setTenant(String tenant) {
        this.tenant = tenant;
    }

    public String getTenant() {
        return tenant;
    }
}

In your application code you can get access to this object in the usual ways, except casting to TenantAwareOAuth2Request, rather than OAuth2Request. Here is an application example:

TenantAwareOAuth2Request request = getOAuth2RequestFromAuthentication();
String clientId = request.getClientId();
String tenant = request.getTenant();
// do something with this information 

And:

public static TenantAwareOAuth2Request getOAuth2RequestFromAuthentication() {
    Authentication authentication = getAuthentication();
    return getTenantAwareOAuth2Request(authentication);
}

private static TenantAwareOAuth2Request getTenantAwareOAuth2Request(Authentication authentication) {
    if (!authentication.getClass().isAssignableFrom(OAuth2Authentication.class)) {
        throw new RuntimeException("unexpected authentication object, expected OAuth2 authentication object");
    }
    return (TenantAwareOAuth2Request) ((OAuth2Authentication) authentication).getOAuth2Request();
}

private static Authentication getAuthentication() {
    SecurityContext securityContext = SecurityContextHolder.getContext();
    return securityContext.getAuthentication();
}

If your resource server is not using Spring Security, there is other libraries to read JWT tokens, and to read the tenant off of these tokens. We have successfully used JJWT for that.

Conclusion

While Spring Security does not make it very easy to add your own properties to JWT tokens, it can certainly be done in an acceptable manner. Having tenant information available in JWT tokens makes these tokens “fully qualified” in a multi-tenant environment, and thus usable without needing additional (tenant) information to be retrieved, when given an access token. This makes it also possible to do multi-tenant-enabled authentication, even on resource servers that aren’t the same component as your authentication server.

You can see how this approach would work for additional properties, on top of just the “tenant” custom property. In fact, make sure that the JWT token contains just enough information so that resource servers can authorize the client without contacting the authorization server.

This post was originally published on Jama’s blog, on June 8th 2016.