Validate the identity token of a user authenticated with AWS Cognito in Java

package com.exampleawsCognito.jwt;

import java.net.MalformedURLException;
import java.net.URL;
import java.security.interfaces.RSAPublicKey;
import java.text.SimpleDateFormat;
import java.util.Calendar;
import java.util.GregorianCalendar;

import com.auth0.jwk.GuavaCachedJwkProvider;
import com.auth0.jwk.Jwk;
import com.auth0.jwk.JwkException;
import com.auth0.jwk.UrlJwkProvider;
import com.auth0.jwt.JWT;
import com.auth0.jwt.JWTVerifier;
import com.auth0.jwt.algorithms.Algorithm;
import com.auth0.jwt.interfaces.DecodedJWT;


/**
 *Handle id tokens obtained by authenticating with Cognito
 *
 * <p>
 *Source that was used as a reference for implementation
 * https://github.com/awslabs/cognito-proxy-rest-service/blob/master/src/main/kotlin/com/budilov/cognito/services/CognitoService.kt
 * </p>
 *
 * @ThreadSafe
 *
 */
public class IdTokenValidator {
	private static final String AWS_REGION = "ap-northeast-1";
	private static final String AWS_COGNITO_USER_POOL_ID = "my_userpool_id";

	private static JWTVerifier verifier = null;
	private static long verifierExpire = -1;
	private static final long VERIFIER_LIVE_MILISEC = 10 * 60 * 1000; //10 minutes
	private static final JWT JWT = new JWT();

	/**
	 *constructor
	 */
	public IdTokenValidator() {
	}

	/**
	 *Validate ID token
	 *
	 * @param idToken ID token to be verified
	 * @ID Token payload if return validation is successful
	 *
	 * @throws InvalidTokenException Authentication failed because the ID Token value is invalid.
	 */
	public DecodedJWT verify(String idToken) throws InvalidTokenException {
		DecodedJWT decodedToken = JWT.decodeJwt(idToken);

		//Make sure it is signed in the cognito user pool
		String iss = decodedToken.getIssuer();
		if (!jwtTokenIssuer().equals(iss)) {
			throw new InvalidTokenException("The issuer of the ID token is not the target system. iss=" + iss, idToken);
		}

		//Make sure that the ID token is used for "ID".
		String tokenUse = decodedToken.getClaim("token_use").asString();
		if (!"id".equals(tokenUse)) {
			throw new InvalidTokenException("The use of the ID token is not an ID. token_use=" + tokenUse, idToken);
		}

		//Check the signature algorithm.
		String alg = decodedToken.getAlgorithm();
		if (!"RS256".equals(decodedToken.getAlgorithm())) {
			throw new InvalidTokenException("The ID token signing algorithm does not support it. alg=" + alg, idToken);
		}

		//Validate payload and signature.
		DecodedJWT decodedJWT = null;
		if ((decodedJWT = tokenVerify(decodedToken)) == null) {
			throw new InvalidTokenException("ID Token verification failed.", idToken);
		}

		return decodedJWT;
	}

	/**
	 *Perform verification using the library of auth0
	 *
	 * @param kid ID Key ID in the token header
	 * @If return is not null, the decoded ID token
	 *
	 * @throws InvalidTokenException Validation failed
	 */
	private DecodedJWT tokenVerify(DecodedJWT jwToken) throws InvalidTokenException {

		try {
			DecodedJWT verified = getVerifier(jwToken.getKeyId()).verify(jwToken);
			return verified;
		} catch (Exception e) {
			throw new InvalidTokenException(e);
		}
	}

	/**
	 *Get an instance of JWTVerifier.
	 *
	 * <p>
	 *JWTVerifier is ver.Since it became thread safe from 3, it will be reused.
	 *However, it should be updated regularly, considering that the RSA key pair used for signing may be updated.
	 * </p>
	 *
	 * @param kid Key ID used for signing
	 *
	 * @return
	 *
	 * @throws MalformedURLException
	 * @throws JwkException
	 */
	private JWTVerifier getVerifier(String kid) throws MalformedURLException, JwkException {
		if (verifier != null && System.currentTimeMillis() < verifierExpire) {
			//If it is within the expiration date, use it as it is
			return verifier;
		}
		synchronized (JWT) {
			//I got the lock, so check it again just in case and then create the instance
			if (verifier != null && System.currentTimeMillis() < verifierExpire) {
				return verifier;
			}

			UrlJwkProvider http = new UrlJwkProvider(new URL(jwksUrl()));
			GuavaCachedJwkProvider provider = new GuavaCachedJwkProvider(http);
			Jwk jwk = provider.get(kid);

			Algorithm algorithm = Algorithm.RSA256((RSAPublicKey) jwk.getPublicKey(), null);
			verifier = JWT.require(algorithm)
					.withIssuer(jwtTokenIssuer())
					.build();

			//Extend the life of JWTVerifier
			verifierExpire = System.currentTimeMillis() + VERIFIER_LIVE_MILISEC;

			Calendar expire = GregorianCalendar.getInstance();
			expire.setTimeInMillis(verifierExpire);
			Logger.info("You have created an instance of JWTVerifier. Deadline"
					+ new SimpleDateFormat("yyyy-MM-dd HH:mm:ss.SSS").format(expire.getTime()));

		}

		return verifier;

	}

	/**
	 *Get the issuer of the ID token
	 *
	 * @return
	 */
	private String jwtTokenIssuer() {
		return String.format("https://cognito-idp.%s.amazonaws.com/%s", AWS_REGION, AWS_COGNITO_USER_POOL_ID);
	}

	/**
	 *JSON web token(JWT)Get the URL of the set.
	 *
	 * @return
	 */
	private String jwksUrl() {
		return jwtTokenIssuer() + "/.well-known/jwks.json";
	}

}

public class InvalidTokenException extends Exception {

    public InvalidTokenException(String message, String idToken) {
        super(message + " token is " + idToken);
    }

    public InvalidTokenException(Throwable e) {
        super(e);
    }
}

    try{
        //TODO IdTokenValidator is thread-safe, so make it static
        DecodedJWT payload = new IdTokenValidator().verify(idToken);
        String cognitoUserName = payload.getClaim("cognito:username").asString();

        //Required processing

    }catch (InvalidTokenException e) {
        Logger.error("Identity token validation failed", e);
        badRequest("The value of IdToken is invalid");
    }

Recommended Posts