Notes on cognito JWT token validation

aws-cognito-token-validation.kt

/***
 * AWS Cognito signs the tokens using the RS256 algorithm. Contrary to most common examples (using HMAC + SHA256) that use
 * a shared secret, the RS256 uses assymetric crytography, so in order to validate the JWT we need to obtain the public key
 * that matches the private key used to generate the token signature.
 * Therefore, the critical step is to download the public key from the discovery endpoint.
 * Cognito makes this available as a JWKS, JSON Web Key Set, available at a well known location of your user pool
 * e.g. https://cognito-idp.{region}.amazonaws.com/{user-pool-id}/.well-known/jwks.json
 * Furthermore, the issuer of the token is the user pool (the iss clamin is the user pool URL) and the audience is the 
 * particular client that was configured on the user pool (the aud claim is set to a specific client id of that pool).
 * See: https://docs.aws.amazon.com/cognito/latest/developerguide/user-pool-settings-client-apps.html
 *
 * Steps: 
 * 1 - Retrieve the JWKS from the Discovery endpoint;
 * 2 - Select the correct key based on the kid property in your JWT
 * 3 - Use that key in order to validate the token's signature
 * 4 - Proceed with the regular claim validations that you'd do for a JWT (e.g. issuer, aud, expiry, etc)
 *
 * Dependencies:
 * - jwt handling library: com.auth0:java-jwt:3.9.0
 * - jwks handling library:  com.auth0:jwks-rsa:0.9.0
*/

package com.example

import com.auth0.jwk.GuavaCachedJwkProvider
import com.auth0.jwk.JwkProvider
import com.auth0.jwk.UrlJwkProvider
import com.auth0.jwt.JWT
import com.auth0.jwt.JWTVerifier
import com.auth0.jwt.algorithms.Algorithm
import org.junit.Test
import java.security.interfaces.RSAPublicKey

fun makeJwtVerifier(issuer: String, audience: String, algorithm: Algorithm): JWTVerifier = JWT
    .require(algorithm)
    .withAudience(audience)
    .withIssuer(issuer)
    .build()


class TestJwtVerifier {
    @Test
    fun `A cognito ID token can be validated`(){
        val cognitoToken = "TODO: REPLACE WITH AN ACTUAL TOKEN"
        val provider: JwkProvider = GuavaCachedJwkProvider(UrlJwkProvider("https://cognito-idp.{region}.amazonaws.com/{user-pool-id}/"))
        val jwt = JWT.decode(cognitoToken)
        val jwk = provider[jwt.keyId]
        val algorithm = Algorithm.RSA256(jwk.publicKey as RSAPublicKey, null)
        val issuer = "https://cognito-idp.{region}.amazonaws.com/{user-pool-id}/"
        val aud = "TODO: REPLACE WITH CLIENT ID"
        
        try{
           makeJwtVerifier(issuer, aud, algorithm).verify(cognitoToken)
        } catch (ex: JWTVerificationException) {
            fail("Token should be valid. Got JWTVerificationException instead")
        }
    }
}