Preload JWKS (#213)

Author: ottokruse

example-serverless-app-reuse/reuse-auth-only.yaml

@@ -32,7 +32,7 @@ Parameters:
   SemanticVersion:
     Type: String
     Description: Semantic version of the back end
-    Default: 2.1.2
+    Default: 2.1.3
 
   HttpHeaders:
     Type: String

example-serverless-app-reuse/reuse-complete-cdk.ts

@@ -19,7 +19,7 @@ const authAtEdge = new sam.CfnApplication(stack, "AuthorizationAtEdge", {
   location: {
     applicationId:
       "arn:aws:serverlessrepo:us-east-1:520945424137:applications/cloudfront-authorization-at-edge",
-    semanticVersion: "2.1.2",
+    semanticVersion: "2.1.3",
   },
   parameters: {
     EmailAddress: "[email protected]",

example-serverless-app-reuse/reuse-complete.yaml

@@ -12,7 +12,7 @@ Resources:
     Properties:
       Location:
         ApplicationId: arn:aws:serverlessrepo:us-east-1:520945424137:applications/cloudfront-authorization-at-edge
-        SemanticVersion: 2.1.2
+        SemanticVersion: 2.1.3
   AlanTuring:
     Type: AWS::Cognito::UserPoolUser
     Properties:

example-serverless-app-reuse/reuse-with-existing-user-pool.yaml

@@ -75,7 +75,7 @@ Resources:
     Properties:
       Location:
         ApplicationId: arn:aws:serverlessrepo:us-east-1:520945424137:applications/cloudfront-authorization-at-edge
-        SemanticVersion: 2.1.2
+        SemanticVersion: 2.1.3
       Parameters:
         UserPoolArn: !GetAtt UserPool.Arn
         UserPoolClientId: !Ref UserPoolClient

src/cfn-custom-resources/fetch-jwks/cfn-response.ts

@@ -0,0 +1,46 @@
+// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: MIT-0
+
+import { request } from "https";
+
+export enum Status {
+  "SUCCESS" = "SUCCESS",
+  "FAILED" = "FAILED",
+}
+
+export async function sendCfnResponse(props: {
+  event: {
+    StackId: string;
+    RequestId: string;
+    LogicalResourceId: string;
+    ResponseURL: string;
+  };
+  status: Status;
+  reason?: string;
+  data?: {
+    [key: string]: string;
+  };
+  physicalResourceId?: string;
+}) {
+  const response = {
+    Status: props.status,
+    Reason: props.reason?.toString() || "See CloudWatch logs",
+    PhysicalResourceId: props.physicalResourceId || "no-explicit-id",
+    StackId: props.event.StackId,
+    RequestId: props.event.RequestId,
+    LogicalResourceId: props.event.LogicalResourceId,
+    Data: props.data || {},
+  };
+
+  await new Promise<void>((resolve, reject) => {
+    const options = {
+      method: "PUT",
+      headers: { "content-type": "" },
+    };
+    request(props.event.ResponseURL, options)
+      .on("error", (err) => {
+        reject(err);
+      })
+      .end(JSON.stringify(response), "utf8", resolve);
+  });
+}

src/cfn-custom-resources/fetch-jwks/https.ts

@@ -0,0 +1,45 @@
+// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: MIT-0
+
+import { request } from "https";
+import { Writable, pipeline } from "stream";
+
+export async function fetch(uri: string) {
+  return new Promise<Buffer>((resolve, reject) => {
+    const req = request(uri, (res) =>
+      pipeline([res, collectBuffer(resolve)], done)
+    );
+
+    function done(error?: Error | null) {
+      if (!error) return;
+      req.destroy(error);
+      reject(error);
+    }
+
+    req.on("error", done);
+
+    req.end();
+  });
+}
+
+const collectBuffer = (callback: (collectedBuffer: Buffer) => void) => {
+  const chunks = [] as Buffer[];
+  return new Writable({
+    write: (chunk, _encoding, done) => {
+      try {
+        chunks.push(chunk);
+        done();
+      } catch (err) {
+        done(err as Error);
+      }
+    },
+    final: (done) => {
+      try {
+        callback(Buffer.concat(chunks));
+        done();
+      } catch (err) {
+        done(err as Error);
+      }
+    },
+  });
+};

src/cfn-custom-resources/fetch-jwks/index.ts

@@ -0,0 +1,72 @@
+// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: MIT-0
+
+import {
+  CloudFormationCustomResourceHandler,
+  CloudFormationCustomResourceDeleteEvent,
+  CloudFormationCustomResourceUpdateEvent,
+} from "aws-lambda";
+import { sendCfnResponse, Status } from "./cfn-response";
+import { fetch } from "./https";
+
+async function fetchJwks(
+  action: "Create" | "Update" | "Delete",
+  userPoolArn: string,
+  physicalResourceId?: string
+) {
+  if (action === "Delete") {
+    // Deletes aren't executed
+    return { physicalResourceId: physicalResourceId!, Data: {} };
+  }
+  console.log(`Fetching JWKS for ${userPoolArn}`);
+
+  const match = userPoolArn.match(
+    new RegExp("userpool/(?<region>.+)_(?<userPoolId>.+)$")
+  );
+  if (!match?.groups) {
+    throw new Error("Failed to parse User Pool ARN");
+  }
+  const url = `https://cognito-idp.${match.groups.region}.amazonaws.com/${match.groups.region}_${match.groups.userPoolId}/.well-known/jwks.json`;
+
+  console.log(`Fetching JWKS from ${url}`);
+  const jwks = (await fetch(url)).toString();
+  console.log(`Fetched JWKS: ${jwks}`);
+
+  return {
+    physicalResourceId: userPoolArn,
+    Data: { Jwks: jwks },
+  };
+}
+
+export const handler: CloudFormationCustomResourceHandler = async (event) => {
+  const { ResourceProperties, RequestType } = event;
+
+  const { PhysicalResourceId } = event as
+    | CloudFormationCustomResourceDeleteEvent
+    | CloudFormationCustomResourceUpdateEvent;
+
+  const { UserPoolArn } = ResourceProperties;
+
+  let status = Status.SUCCESS;
+  let physicalResourceId: string | undefined;
+  let data: { [key: string]: any } | undefined;
+  let reason: string | undefined;
+  try {
+    ({ physicalResourceId, Data: data } = await fetchJwks(
+      RequestType,
+      UserPoolArn,
+      PhysicalResourceId
+    ));
+  } catch (err) {
+    console.error(err);
+    status = Status.FAILED;
+    reason = `${err}`;
+  }
+  await sendCfnResponse({
+    event,
+    status,
+    data,
+    physicalResourceId,
+    reason,
+  });
+};

src/cfn-custom-resources/fetch-jwks/package.json

@@ -0,0 +1,11 @@
+{
+  "name": "fetch-jwks",
+  "version": "1.0.0",
+  "description": "",
+  "main": "index.js",
+  "scripts": {
+    "test": "echo \"Error: no test specified\" && exit 1"
+  },
+  "keywords": [],
+  "author": ""
+}

src/lambda-edge/shared/shared.ts

@@ -10,6 +10,7 @@ import { fetch } from "./https";
 import { Agent, RequestOptions } from "https";
 import html from "./error-page/template.html";
 import { CognitoJwtVerifier } from "aws-jwt-verify";
+import { Jwks } from "aws-jwt-verify/jwk";
 export {
   CognitoJwtInvalidGroupError,
   JwtExpiredError,
@@ -75,6 +76,7 @@ interface ConfigFromDiskWithHeaders extends ConfigFromDisk {
 
 interface ConfigFromDiskComplete extends ConfigFromDiskWithHeaders {
   userPoolArn: string;
+  jwks: Jwks;
   clientId: string;
   oauthScopes: string[];
   cognitoAuthDomain: string;
@@ -228,14 +230,21 @@ export function getCompleteConfig(): CompleteConfig {
 export function getConfigWithJwtVerifier() {
   const config = getCompleteConfig();
   const userPoolId = config.userPoolArn.split("/")[1];
+  const jwtVerifier = CognitoJwtVerifier.create({
+    userPoolId,
+    clientId: config.clientId,
+    tokenUse: "id",
+    groups: config.requiredGroup || undefined,
+  });
+
+  // Optimization: load the JWKS (as it was at deploy-time) into the cache.
+  // Then, the JWKS does not need to be fetched at runtime,
+  // as long as only JWTs come by with a kid that is in this cached JWKS:
+  jwtVerifier.cacheJwks(config.jwks);
+
   return {
     ...config,
-    jwtVerifier: CognitoJwtVerifier.create({
-      userPoolId,
-      clientId: config.clientId,
-      tokenUse: "id",
-      groups: config.requiredGroup || undefined,
-    }),
+    jwtVerifier,
   };
 }
 

template.yaml

@@ -27,7 +27,7 @@ Metadata:
         "amplify",
       ]
     HomePageUrl: https://github.com/aws-samples/cloudfront-authorization-at-edge
-    SemanticVersion: 2.1.2
+    SemanticVersion: 2.1.3
     SourceCodeUrl: https://github.com/aws-samples/cloudfront-authorization-at-edge
 
 Parameters:
@@ -150,7 +150,7 @@ Parameters:
   Version:
     Type: String
     Description: "Changing this parameter after initial deployment forces redeployment of Lambda@Edge functions"
-    Default: "2.1.2"
+    Default: "2.1.3"
   LogLevel:
     Type: String
     Description: "Use for development: setting to a value other than none turns on logging at that level. Warning! This will log sensitive data, use for development only"
@@ -444,6 +444,7 @@ Resources:
                 - lambda:UpdateFunctionCode
                 - lambda:UpdateFunctionConfiguration
                 - lambda:TagResource
+                - lambda:ListTags
               Resource:
                 - !Sub "arn:${AWS::Partition}:lambda:us-east-1:${AWS::AccountId}:function:*-CheckAuthHandler-*"
                 - !Sub "arn:${AWS::Partition}:lambda:us-east-1:${AWS::AccountId}:function:*-ParseAuthHandler-*"
@@ -754,6 +755,22 @@ Resources:
                 - !GetAtt UserPool.Arn
                 - !Ref UserPoolArn
 
+  CognitoJwksFetchHandler:
+    Type: AWS::Serverless::Function
+    Properties:
+      CodeUri: src/cfn-custom-resources/fetch-jwks/
+      Handler: index.handler
+
+  FetchedJwks:
+    Type: Custom::FetchedJwks
+    Properties:
+      ServiceToken: !GetAtt CognitoJwksFetchHandler.Arn
+      Version: !Ref Version
+      UserPoolArn: !If
+        - CreateUserPoolAndClient
+        - !GetAtt UserPool.Arn
+        - !Ref UserPoolArn
+
   StaticSite:
     Type: Custom::StaticSite
     Condition: CreateSampleStaticSite
@@ -849,6 +866,7 @@ Resources:
           - >
             {
               "userPoolArn": "${UserPoolArn}",
+              "jwks": ${FetchedJwks.Jwks},
               "clientId": "${ClientId}",
               "clientSecret": "${ClientSecret}",
               "oauthScopes": ${OAuthScopesJsonArray},
@@ -911,6 +929,7 @@ Resources:
           - >
             {
               "userPoolArn": "${UserPoolArn}",
+              "jwks": ${FetchedJwks.Jwks},
               "clientId": "${ClientId}",
               "clientSecret": "${ClientSecret}",
               "oauthScopes": ${OAuthScopesJsonArray},
@@ -1005,6 +1024,7 @@ Resources:
           - >
             {
               "userPoolArn": "${UserPoolArn}",
+              "jwks": ${FetchedJwks.Jwks},              
               "clientId": "${ClientId}",
               "clientSecret": "${ClientSecret}",
               "oauthScopes": ${OAuthScopesJsonArray},
@@ -1067,6 +1087,7 @@ Resources:
           - >
             {
               "userPoolArn": "${UserPoolArn}",
+              "jwks": ${FetchedJwks.Jwks},              
               "clientId": "${ClientId}",
               "clientSecret": "${ClientSecret}",
               "oauthScopes": ${OAuthScopesJsonArray},