April 05, 2020

AWS Cognito & JWT

Recently I had to design and implement a solution that used a third-party user management system for authentication. I decided to use Amazon’s Cognito service, more specifically the User Pool aspect. The third-party service was able to work with SAML, and so does Cognito. Cognito’s output that you use is a JWT object. The backend system is written using Java 8 and Spring Framework and Spring Security.

Using Spring Boot parent POM:

<parent>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-parent</artifactId>
    <version>2.1.5.RELEASE</version>
    <relativePath/>
</parent>

I have a great deal of experience building systems using Spring Security because of its extremely adaptable security model. Initially I had attempted to use the Auth0 library as they had modules that were ready for Spring Security. For several reasons I need to make modifications to those libraries and they were designed in a way that made it impossible to simply extend those libraries. I ended up taking some of that code and just copying it and pasting it into my base classes and work from there.

The two dependencies I am using for managing JWT objects is:

<dependencies>
    <dependency>
        <groupId>com.auth0</groupId>
        <artifactId>java-jwt</artifactId>
    </dependency>
    <dependency>
        <groupId>com.auth0</groupId>
        <artifactId>jwks-rsa</artifactId>
    </dependency>
</dependencies>

If you are interested you can probably use:

<dependency>
    <groupId>com.auth0</groupId>
    <artifactId>auth0-spring-security-api</artifactId>
    <version>1.3.0</version>
</dependency>

There are several things I want to mention with my security model that required a lot of tweaks:

We are using the user groups that are passed along in the JWT object in cognito:groups key.
Needed to support a custom “granting” system that allows certain calls to “grant” itself certain roles or permissions for the duration of the call without requiring the user to possess those roles/permissions.
Needed to support service tokens, external services that access the API and will not authenticate by means of a JWT. The list of authorities for each service would be pre-defined.

«««< HEAD

Support for multiple JWT authorization sources

=======

Support for multiple JWT authorization sources

staging

Part of what makes Spring Security such a great system is that almost everything can be configured and tweaked, and at the same time like all Spring libraries they are sensible defaults. An AuthenticationProvider is an interface that is very generic and at its core the interface used to actually either authenticate or reject a support Authenatication. There are many layers built-in to Spring Security for common patterns and needs such as the UserDetailsService. Spring Security supports at the same time multiple AuthenticationProviders but for multiple custom AuthenticationProvider you will need to use an AuthenticationManager. For testing purposes I wanted to allow the use of real user credentials from the third-party service as well as credentials that are generated dynamically for the various use-cases as needed. The first that you need to figure out is how to create your own JWT local keystore. That means that you should be able to cryptographically create the token prior to a test and then be able to have the backend verify it and use it correctly. With limited time and resources if I am writing tests for the front-end they are generally going to be end-to-end tests, not unit tests. That is why I wanted to use real tokens and make sure that the correct roles work at the right times and the opposite as well.

I created my own version of JwtAuthenticationProvider that takes in a Set<JwtAuthorizer. I made up that class but it is the constructor parameters that the JwtAuthenticationProvider uses.

public static class JwtAuthorizer {
    private final JwkProvider jwkProvider;
    private final byte[] secret;
    private final String issuer;
    private final String audience;

    public JwtAuthorizer(JwkProvider jwkProvider, String issuer, String audience, byte[] secret) {
        this.jwkProvider = jwkProvider;
        this.issuer = issuer;
        this.audience = audience;
        this.secret = secret;
    }

    public JwtAuthorizer(JwkProvider jwkProvider, String issuer) {
        this.jwkProvider = jwkProvider;
        this.issuer = issuer;
        this.audience = null;
        this.secret = null;
    }

    public JwkProvider getJwkProvider() {
        return jwkProvider;
    }

    public String getIssuer() {
        return issuer;
    }

    public String getAudience() {
        return audience;
    }

    public byte[] getSecret() {
        return secret;
    }
}

Here is my slightly modified version of the JwtAuthenticationProvider. The original source is available here.

public class JwtAuthenticationProvider implements AuthenticationProvider {

    private static Logger logger = LoggerFactory.getLogger(JwtAuthenticationProvider.class);

    private final Set<JwtAuthorizer> authorizers = new HashSet();

    private long leeway = 0;

    public JwtAuthenticationProvider(Set<JwtAuthorizer> authorizers) {
        this.authorizers.addAll(authorizers);
    }

    public JwtAuthenticationProvider(JwkProvider jwkProvider, String issuer, String audience, byte[] secret) {
        this.authorizers.add(new JwtAuthorizer(jwkProvider, issuer, audience, secret));
    }

    public JwtAuthenticationProvider(JwkProvider jwkProvider, String issuer, String audience) {
        this.authorizers.add(new JwtAuthorizer(jwkProvider, issuer, audience, null));
    }

    @Override
    public boolean supports(Class<?> authentication) {
        return JwtAuthentication.class.isAssignableFrom(authentication);
    }

    @Override
    public Authentication authenticate(Authentication authentication) throws AuthenticationException {
        if (!supports(authentication.getClass())) {
            return null;
        }

        JwtAuthentication jwt = (JwtAuthentication) authentication;
        for(JwtAuthorizer authorizer : authorizers) {
            try {
                final Authentication jwtAuth = jwt.verify(jwtVerifier(jwt));
                logger.info("Authenticated with jwt with scopes {}", jwtAuth.getAuthorities());
                return jwtAuth;
            } catch (Exception e) {
                // Not needed, just to catch.
                logger.warn("An authentication exception was thrown: {}", e.getMessage());
            }
        }

        throw new BadCredentialsException("Not a valid token");
    }

    /**
     * Allow a leeway to use on the JWT verification.
     *
     * @param leeway the leeway value to use expressed in seconds.
     * @return this same provider instance to chain calls.
     */
    @SuppressWarnings("unused")
    public JwtAuthenticationProvider withJwtVerifierLeeway(long leeway) {
        this.leeway = leeway;
        return this;
    }

    private JWTVerifier jwtVerifier(JwtAuthentication authentication) throws AuthenticationException {
        if (secret != null) {
            return providerForHS256(secret, issuer, audience, leeway);
        }
        final String kid = authentication.getKeyId();
        if (kid == null) {
            throw new BadCredentialsException("No kid found in jwt");
        }
        if (jwkProvider == null) {
            throw new AuthenticationServiceException("Missing jwk provider");
        }
        try {
            final Jwk jwk = jwkProvider.get(kid);
            return providerForRS256((RSAPublicKey) jwk.getPublicKey(), issuer, audience, leeway);
        } catch (SigningKeyNotFoundException e) {
            throw new AuthenticationServiceException("Could not retrieve jwks from issuer", e);
        } catch (InvalidPublicKeyException e) {
            throw new AuthenticationServiceException("Could not retrieve public key from issuer", e);
        } catch (JwkException e) {
            throw new AuthenticationServiceException("Cannot authenticate with jwt", e);
        }
    }

    private static JWTVerifier providerForRS256(RSAPublicKey publicKey, String issuer, String audience, long leeway) {
        return JWT.require(Algorithm.RSA256(publicKey, null))
                .withIssuer(issuer)
                .withAudience(audience)
                .acceptLeeway(leeway)
                .build();
    }

    private static JWTVerifier providerForHS256(byte[] secret, String issuer, String audience, long leeway) {
        return JWT.require(Algorithm.HMAC256(secret))
                .withIssuer(issuer)
                .withAudience(audience)
                .acceptLeeway(leeway)
                .build();
    }
}

The only real change I’ve made here was iterating through the configured list of JwtAuthorizers. There is an alternative approach that I could have taken for this implementation. Spring Security uses the boolean supports(Class<?> authentication) method to determine if the AuthenticationProvider is appropriate for the Authentication supplied. Unfortunately it doesn’t allow for a mechanism to indicate that the provider is the right type but the wrong authorizer. Alternatively, you might be able to check if the Authentication object is an instance of JwtAuthentication and that it’s kid matches the authorizer’s. I prefer to use the supports method for determining the type match, not getting into the more nitty-gritty of the which authorizer. I like the idea of only having a single JwtAuthenticationProvider and making the JwtAuthorizer programmatic. This would allow for runtime configuration and as many authorizer instances as desired.

Now we need to configure support for our locally served jwks.json file. My system configuration uses docker and since the backend is setup first and then then the nginx to host the front-end. I needed to make sure that I could run tests if I wanted to using JWT that didn’t depend on the front-end being up, especially in an environment like CodeBuild. I decided to store the jwks.json as a classpath resource. This is a very simple class that implements the JwkProvider interface.

public class LocalJwkProvider implements JwkProvider {
	private InputStream input;
	private Map<String, Object> keys;

	public LocalJwkProvider(InputStream input) {
		this.input = input;
		try {
			this.keys = getJwks();
		} catch (Exception e) {

		}
	}

	private Map<String, Object> getJwks() throws SigningKeyNotFoundException {
		try {
			final JsonFactory factory = new JsonFactory();
			final JsonParser parser = factory.createParser(this.input);
			final TypeReference<Map<String, Object>> typeReference = new TypeReference<Map<String, Object>>() {};
			return new ObjectMapper().reader().readValue(parser, typeReference);
		} catch (IOException e) {
			throw new SigningKeyNotFoundException("Problem with parsing...", e);
		}
	}

	static Jwk fromValues(Map<String, Object> map) {
		Map<String, Object> values = Maps.newHashMap(map);
		String kid = (String) values.remove("kid");
		String kty = (String) values.remove("kty");
		String alg = (String) values.remove("alg");
		String use = (String) values.remove("use");
		Object keyOps = values.remove("key_ops");
		String x5u = (String) values.remove("x5u");
		List<String> x5c = (List<String>) values.remove("x5c");
		String x5t = (String) values.remove("x5t");
		if (kty == null) {
			throw new IllegalArgumentException("Attributes " + map + " are not from a valid jwk");
		}
		if(keyOps instanceof String) {
			return new Jwk(kid, kty, alg, use, (String) keyOps, x5u, x5c, x5t, values);
		} else {
			return new Jwk(kid, kty, alg, use, (List<String>) keyOps, x5u, x5c, x5t, values);
		}
	}

	private List<Jwk> getAll() throws SigningKeyNotFoundException {
		List<Jwk> jwks = Lists.newArrayList();
		@SuppressWarnings("unchecked") 
        final List<Map<String, Object>> keys = (List<Map<String, Object>>) this.keys.get("keys");

		if (keys == null || keys.isEmpty()) {
			throw new SigningKeyNotFoundException("No keys found", null);
		}

		try {
			for (Map<String, Object> values : keys) {
				jwks.add(fromValues(values));
			}
		} catch (IllegalArgumentException e) {
			throw new SigningKeyNotFoundException("Failed to parse jwk from json", e);
		}
		return jwks;
	}

	@Override
	public Jwk get(String keyId) throws JwkException {
		final List<Jwk> jwks = getAll();
		if (keyId == null && jwks.size() == 1) {
			return jwks.get(0);
		}
		if (keyId != null) {
			for (Jwk jwk : jwks) {
				if (keyId.equals(jwk.getId())) {
					return jwk;
				}
			}
		}
		throw new SigningKeyNotFoundException("", null);
	}
}

In order to use this I just added a bit to my existing SecurityConfig:

@Configuration
@EnableWebSecurity
public class SecurityConfig extends WebSecurityConfigurerAdapter {

	@Value("${jwt.url}")
	private String jwtUrl;

	@Autowired
	private Environment environment;

	@Bean
	public JwtAuthenticationProvider jwtAuthenticationProvider() throws IOException {
		Set<JwtAuthenticationProvider.JwtAuthorizer> authorizers = new HashSet();
		final JwkProvider jwkProvider = new UrlJwkProvider(URI.create(jwtUrl+"/.well-known/jwks.json").toURL());
		authorizers.add(new JwtAuthorizer(jwkProvider, jwtUrl));

        // We don't want to have this enabled normally
        // We need to inject the different accepted JWT authorization source into our JwtAuthenticationProvider
        // The reason 
		if(ArrayUtils.contains(environment.getActiveProfiles(), "test")) {
			ClassPathResource r = new ClassPathResource("/jwks.json");
			byte[] bytes = IOUtils.toByteArray(r.getInputStream());
			JwkProvider localProvider = new LocalJwkProvider(new ByteArrayInputStream(bytes));
			authorizers.add(
				new JwtAuthenticationProvider.JwtAuthorizer(localProvider, "http://localhost")
			);
		}

		return new JwtAuthenticationProvider(authorizers);
	}

	@Bean
	public AuthenticationManager authenticationManager() throws IOException {
		return new ProviderManager(
			Lists.newArrayList(
				jwtAuthenticationProvider(), serviceAuthenticationProvider()
			)
		);
	}

    // Other code...
}

The authenticationManager() allows you to have multiple AuthenticationProvider instances for different types of Authentication you might need. This configuration enables you to keep your production authentication as a valid means of authenticating if you want. While this isn’t common, at times I will be working with a local copy of production data and want to try something out myself locally with real credentials but no negative system impact. This is especially useful when debugging an issue and trying to replicate the issue. With the way I have designed it as far as the sytem

To begin this you will need to generate your jwks.json file. There is a nice free site available that will help you do this without any fuss. Once you have that jwks.json file
The library you are going to want to use on the front-end is jsrsasign. Here is a little JavaScript library to generate a token.

const $ = require('jquery');
const _ = require('lodash');
const KJUR = require('jsrsasign').KJUR;
const RSAKey = require('jsrsasign').RSAKey;


function random(num) {
    var nums = _.range(48, 57);
    var lettersU = _.range(65, 90);
    var lettersL = _.range(97, 122);

    var all = _.union(nums, lettersL, lettersU);

    var str = '';
    for(i=0; i<= num; i++) {
        var index = _.random(0, all.length-1);
        var val = all[index];
        str += String.fromCodePoint(val)
    }
    return str;
}

function generateToken(payload, key) {
    // Replace your key id
    var header = {alg: 'RS256', kid: 'KEY_ID'};
    var sHeader = JSON.stringify(header);
    var sPayload = JSON.stringify(payload);
    var sJWT = KJUR.jws.JWS.sign("RS256", sHeader, sPayload, key);

    return sJWT;
}

function createTestCredentials(username, groups, key, cb) {
    if(typeof username == 'undefined') {
        username = 'test';
    }

    if(typeof key == 'undefined') {
        key = getPem();
    }

    var rsa = new RSAKey();
    rsa.readPrivateKeyFromPEMString(key);

    var tNow = KJUR.jws.IntDate.get('now');
    var tEnd = KJUR.jws.IntDate.get('now + 1day');

    var payload = {
        'iss': window.location.protocol+"://"+window.location.host,
        'iat': tNow,
        'username': username,
        'jti': random(50),
        'exp': tEnd,
        'cognito:groups': groups
    };

    var res = generateToken(payload, rsa);
    console.log('Test token generated');

    if(typeof cb !== 'undefined') {
        cb(res);
    }

    return res;
}

function storeTestCredentials(token) {
    localStorage.setItem('token', token);
}

function createAndStoreTestCredentials(username, groups, key) {
    return createTestCredentials(username, groups, key, storeTestCredentials);
}

function getPem() {
    var url;
    if(_.isEqual(window.location.host, 'localhost')) {
        url = 'http://localhost/test.pem';
    } else {
        url = window.location.protocol+'://'+window.location.host+'/test.pem';
    }

    return $.ajax({
        type: "GET",
        url: url,
        async: false
    }).responseText;
}

if (typeof module !== 'undefined' && typeof module.exports !== 'undefined') {
    module.exports = {
        createTestCredentials: createTestCredentials,
        createAndStoreTestCredentials: createAndStoreTestCredentials,
        getPem: getPem
    };
}
window.testAuth = module.exports;

So you will need to serve the test.pem file locally. If you used the “mkjwk” site I listed above you will need to extract the private key from the JSON and put it as a file. The net result of using this will produce a JWT object that you are signing yourself but in all other ways will work with cognito or really anything else that generates a JWT token.

I have two different ways that I use this. The first way is just “Fake Login” which just creates a token with the ADMIN role and “test” as the username. I use this when I really just need to get into the system and either don’t care or don’t want to be troubled with thinking about the exact roles I need at the moment. The other approach is where I specify the exact roles or groups. Either way its the the same here.

«««< HEAD

Support for granting roles and permission dynamically

=======

Support for granting roles and permission dynamically

staging This is the more complex functionality that requires both knowledge of Spring Security as well how Spring Framework operates. I’m going to give you the outline of the approach now. Quite simply we want to be able to give users access to certain roles and permission but not all the time. Really, we want them to be able to perform operation A, B, and C. Normally knowing that operation A requires permission X, Y, and Z any attempt to invoke operation A without the requisite permission would yield in a 403. Operation A is already protected by requiring permission W. We don’t want the user to have general permissions X, Y, and Z on a normal basis, but we do want them to be able to invoke operation A. What you would have to do is grant the user permissions X, Y, and Z prior to the execution of operation A and remove them once the invocation has completed.

Let’s see the very simple annotation I created:

@Target({ElementType.METHOD, ElementType.TYPE})
@Retention(RetentionPolicy.RUNTIME)
public @interface Grants {
	Permission[] value() default {};
	Role[] roles() default {};
	String name() default "";
}

What I have done with this is created a simple @Aspect that has a @Before and @After JointPoint advice. Simply enough I created a base Authentication interface called GrantingAuthentication which extends the standard Authentication interface.

public interface GrantingAuthentication extends Authentication {
	Collection<GrantedAuthority> getAuthorities();

	void addAuthority(Collection<GrantedAuthority> authorities);
	void removeAuthorities(Collection<Authority> authorities);
}

With this interface you need to modify any of your Authentication implementations to implement these functions. Using this interface your aspect component can simply add the appropriate authorities (roles or permissions) prior to invocation. One thing that you need to be careful of is to keep track of what authorities you added, otherwise you might remove authorities that the user actual has outside of the grant.

«««< HEAD

Service Token

=======

Service Token

staging To handle the service token need there is a very simple and elegant solution. Use the SecurityContextRepository interface. Both the JWT and the service token take use the Authorization header with the value prefixed by Bearer. With that all you need to do is in the SecurityContext loadContext(HttpRequestResponseHolder requestResponseHolder) function obtain the token from the header. Then if the token is one of the defined service tokens then just use SecurityContextHolder.setContext(context) to populate a security context with the appropriate authorities. If the token is not for a service token then use the PreAuthenticatedAuthenticationJsonWebToken.usingToken(token) call. The Auth0 JWT implementation already uses this approach here.

«««< HEAD

Summary

This approach has served my purposes quite well, and I have used several aspects of this approach in many other projects. Reusing good solid architecture is the corner stone of a strong foundation that is essential everywhere and all the time. This resource is in no way exhaustive or intended to be complete. I hope some things here might help someone else with similar needs as myself or my organziation. As always. comments and questions are always welcome. Enjoy, and make beautiful code. =======

Summary

staging