express.js middleware to validate a AWS Cognito / Amplify Token

auth.ts

import { Router } from "express";
import jwt from "jsonwebtoken";
import jwkToPem from "jwk-to-pem";
import * as Axios from 'axios';

interface PublicKeys {
  keys: PublicKey[];
}
interface PublicKey {
  alg: string;
  e: string;
  kid: string;
  kty: string;
  n: string;
  use: string;
}
interface PublicKeyMeta {
  instance: PublicKey;
  pem: string;
}
interface MapOfKidToPublicKey {
  [key: string]: PublicKeyMeta;
}
let cacheKeys: MapOfKidToPublicKey | undefined;

// TODO - abstract pool id
const cognitoPoolId = process.env.COGNITO_POOL_ID || '';
const cognitoIssuer = `https://cognito-idp.us-east-1.amazonaws.com/${cognitoPoolId}`;

if (!cognitoPoolId) {
  throw new Error('env var required for cognito pool');
}

/**
 * @returns - returns a map of AWS Cognito public keys by kid (key ID). this public key can be used to verify the signature of a jwt token
 */
const getPublicKeys = async (): Promise<MapOfKidToPublicKey> => {
  if (!cacheKeys) {
    const url = `${cognitoIssuer}/.well-known/jwks.json`;
    const publicKeys = await Axios.default.get<PublicKeys>(url);
    cacheKeys = publicKeys.data.keys.reduce((agg, current) => {
      const pem = jwkToPem(current);
      agg[current.kid] = { instance: current, pem };
      return agg;
    }, {} as MapOfKidToPublicKey);
    return cacheKeys;
  } else {
    return cacheKeys;
  }
};

/**
 *
 * Resources
 * https://cognito-idp.us-east-1.amazonaws.com/{cognitoPoolId}.well-known/jwks.json
 *
 * https://tools.ietf.org/html/rfc7517
 *
 * @param - kid - Key ID of public key used to sign JWT
 * @returns - return a JSON Web Key (JWK) for kid
 */
const getJwkByKid = async (kid: string): Promise<PublicKey> => {
  const keys = await getPublicKeys();
  const publicKey: PublicKeyMeta = keys[kid];

  // if the public key is missing reload cache and try one more time
  // https://forums.aws.amazon.com/message.jspa?messageID=747599
  if (!publicKey) {
    cacheKeys = undefined;
    const keys2 = await getPublicKeys();
    const publicKey2: PublicKeyMeta = keys2[kid];
    if (!publicKey2) {
      return null;
    }
    return publicKey2.instance;
  }
  return publicKey.instance;
}

/**
 *
 * Map an auth token onto request
 * @param req - request
 * @param res - response
 * @param next - request callback
 */
const getAuthToken = (req, res, next) => {
  if (
    req.headers.authorization &&
    req.headers.authorization.split(' ')[0] === 'Bearer'
  ) {
    req.authToken = req.headers.authorization.split(' ')[1];
  } else {
    req.authToken = null;
  }
  next();
};

/**
 *
 * 
 * @param req - Express Request. it is expected that req.authToken has been populated
 */
const validateToken = async (req) => {
  const { authToken } = req;
  if(authToken === null){
    throw new Error("auth token is not present in request");
  }
  // https://github.com/awslabs/aws-support-tools/tree/master/Cognito/decode-verify-jwt
  const jwtDecoded = jwt.decode(authToken, { complete: true });
  const jwk = await getJwkByKid(jwtDecoded.header.kid);
  if (jwk === null) {
    throw new Error("unable to validate token");
  }
  const pem = jwkToPem(jwk);
  await jwt.verify(authToken, pem, { algorithms: ['RS256'] }, (err, decoded) => {
    if (decoded.token_use !== 'access' || decoded.iss !== cognitoIssuer) {
      throw new Error("unauthorized user");
    }
  });
}

/**
 *
 * Map an auth toekn onto a request and validate it. we only need to validate player voting API calls
 *
 * @param req - express request
 * @param res - express response
 * @param next - request callback
 */
const verifyAuthToken = (req, res, next) => {
  getAuthToken(req, res, async () => {
    try {
      await validateToken(req);
      return next();
    } catch (e) {
      return res
        .status(401)
        .send({ error: 'You are not authorized to make this request' });
    }
  });
};

export const checkIfAuthenticated = (router: Router) => {
  router.use(verifyAuthToken);
};