How to use JWT tokens for authorization with AWS KMS in NodeJs Lambdas

Let’s look at how to sign and verify JWT tokens with AWS Key Management Service (KMS) keys in NodeJs18-based Lambda functions.

This will save you a couple of days as both tutorials on what data types are needed (UInt8 Arrays? Buffers?) and official documentation are inconclusive.

The article will provide two lambda functions used respectively for the JWT token generation and token validation via authoriser.

Create the key

To sign and verify JWT tokens you have to create an asymmetric KMS key: Creating asymmetric KMS keys

In this context, I used the following:

• Key Usage: Sign and Verify

• KeySpec: RSA_2048

A quick recap on JWT tokens

Always refer to RFC when looking for a proper definition: RFC 7519: JSON Web Token (JWT)

In a nutshell, a JWT token is a period-separated concatenation of 3 bases64url encoded components:

• Header, with info algorithm and signed key, among others.

• Payload, containing exchanged information as a set of claims.

• Signature, to make sure the token is trustworthy.

Token Generation

The following lambda code generates a JWT token with no JWT third-party libraries (Jose, jsonwebtoken, etc…).

import { KMSClient, SignCommand } from ‘@aws-sdk/client-kms’;
import base64url from ‘base64url’;

const kmsClient = new KMSClient({
    region: "eu-central-1"
});

export const generateTokenLambdaHandler = async (event) => {
    const kmsKeyId = process.env.KMS_KEY_ID;

    const token = await sign(kmsKeyId)

    return {
        "isBase64Encoded": false,
        "statusCode": 200,
        "body": JSON.stringify(token)
    }
}

async function sign(kmsKeyId) {

    // The JOSE header
    const headers = {
        "alg": "RS256",
        "typ": "JWT",
        "kid": kmsKeyId
    }

    // The custom payload you wish to send
    let payload = {
        "user_name": "ala"
    }

    payload.iat = Math.floor(Date.now() / 1000);

    const tomorrow = new Date()
    tomorrow.setDate(tomorrow.getDate() + 1)
    payload.exp = Math.floor(tomorrow.getTime() / 1000);

    // Add more claims if you need them

    let token_components = {
        header: base64url(JSON.stringify(headers)),
        payload: base64url(JSON.stringify(payload)),
    };

    let message = Buffer.from(token_components.header + "." + token_components.payload)

    const input = {
        KeyId: kmsKeyId,
        Message: message,
        SigningAlgorithm: "RSASSA_PKCS1_V1_5_SHA_256",
        MessageType: 'RAW'
    };
    const command = new SignCommand(input);
    const res = await kmsClient.send(command);

    token_components.signature = Buffer.from(res.Signature).toString('base64');

    return token_components.header + "." + token_components.payload + "." + token_components.signature;
}

Code should be self explanatory, but here my main regards:

• Base64url encode the 3 components

• Use KMS’ SignCommand with proper SigningAlgorithm

• Do not base64 url encode the signature, but just base64 it!

Token verification

This lambda will verify that token is correctly signed with same KMS key provided in the signature.

import { KMSClient, VerifyCommand } from ‘@aws-sdk/client-kms’;

const kmsClient = new KMSClient({
    region: "eu-central-1"
});

export const authorizerLambdaHandler = async (event) => {
    const kmsKeyId = process.env.KMS_KEY_ID;
   
    let token = event.authorizationToken;
    
   
    if (!token || !token.startsWith("Bearer")) {
       // Handle error accordingly
    }
    
    token = token.replace(/^Bearer\s+/, "");
    
    const [headerBase64, payloadBase64, signatureBase64] = token.split(".");

    const header = Buffer.from(headerBase64, ‘base64').toString()
    const payload = Buffer.from(payloadBase64, ‘base64').toString()
    console.log(header + " " + payload)
    //validate header and payload as you wish

    // The evil line that you need
    const signatureToVerify = Uint8Array.from(Buffer.from(signatureBase64, ‘base64'))

    const input = {
        KeyId: kmsKeyId,
        Message: Buffer.from(headerBase64 + "." + payloadBase64),
        MessageType: "RAW",
        Signature: Buffer.from(signatureToVerify),
        SigningAlgorithm: "RSASSA_PKCS1_V1_5_SHA_256"
    };

    const command = new VerifyCommand(input);
    const response = await kmsClient.send(command);

    if (!response.SignatureValid) {
        return generatePolicy(‘Deny’, event.methodArn)
    }

    return generatePolicy(‘Allow’, event.methodArn)
}

// Help function to generate an IAM policy
const generatePolicy = (effect, resource) => {
    const authResponse = {};
    const principalId = "user"; // use proper
    authResponse.principalId = principalId;
    if (effect && resource) {
        let policyDocument = {};
        policyDocument.Version = '2012-10-17';
        policyDocument.Statement = [];
        let statementOne = {};
        statementOne.Action = 'execute-api:Invoke';
        statementOne.Effect = effect;
        statementOne.Resource = resource;
        policyDocument.Statement[0] = statementOne;
        authResponse.policyDocument = policyDocument;
    }

    return authResponse;
}

Hope this is helpful!