Fixes for cdk-cloudfront-authorization

packages/cdk-cloudfront-authorization/src/__tests__/__snapshots__/authorizations.test.ts.snap

@@ -340,7 +340,8 @@ Object {
               Object {
                 "Ref": "SpaAuthorizationSecretGeneratorCustomResource7B87C244",
               },
-              "\\"
+              "\\",
+  \\"httpHeaders\\": []
 }",
             ],
           ],
@@ -419,7 +420,8 @@ Object {
               Object {
                 "Ref": "SpaAuthorizationSecretGeneratorCustomResource7B87C244",
               },
-              "\\"
+              "\\",
+  \\"httpHeaders\\": []
 }",
             ],
           ],
@@ -498,7 +500,8 @@ Object {
               Object {
                 "Ref": "SpaAuthorizationSecretGeneratorCustomResource7B87C244",
               },
-              "\\"
+              "\\",
+  \\"httpHeaders\\": []
 }",
             ],
           ],
@@ -577,7 +580,8 @@ Object {
               Object {
                 "Ref": "SpaAuthorizationSecretGeneratorCustomResource7B87C244",
               },
-              "\\"
+              "\\",
+  \\"httpHeaders\\": []
 }",
             ],
           ],
@@ -1325,7 +1329,8 @@ Object {
                   "UserPoolClient.ClientSecret",
                 ],
               },
-              "\\"
+              "\\",
+  \\"httpHeaders\\": []
 }",
             ],
           ],
@@ -1412,7 +1417,8 @@ Object {
                   "UserPoolClient.ClientSecret",
                 ],
               },
-              "\\"
+              "\\",
+  \\"httpHeaders\\": []
 }",
             ],
           ],
@@ -1499,7 +1505,8 @@ Object {
                   "UserPoolClient.ClientSecret",
                 ],
               },
-              "\\"
+              "\\",
+  \\"httpHeaders\\": []
 }",
             ],
           ],
@@ -1586,7 +1593,8 @@ Object {
                   "UserPoolClient.ClientSecret",
                 ],
               },
-              "\\"
+              "\\",
+  \\"httpHeaders\\": []
 }",
             ],
           ],

packages/cdk-cloudfront-authorization/src/__tests__/__snapshots__/distributions.test.ts.snap

@@ -393,7 +393,8 @@ Object {
               Object {
                 "Ref": "SpaAuthorizationSecretGeneratorCustomResource7B87C244",
               },
-              "\\"
+              "\\",
+  \\"httpHeaders\\": []
 }",
             ],
           ],
@@ -472,7 +473,8 @@ Object {
               Object {
                 "Ref": "SpaAuthorizationSecretGeneratorCustomResource7B87C244",
               },
-              "\\"
+              "\\",
+  \\"httpHeaders\\": []
 }",
             ],
           ],
@@ -551,7 +553,8 @@ Object {
               Object {
                 "Ref": "SpaAuthorizationSecretGeneratorCustomResource7B87C244",
               },
-              "\\"
+              "\\",
+  \\"httpHeaders\\": []
 }",
             ],
           ],
@@ -630,7 +633,8 @@ Object {
               Object {
                 "Ref": "SpaAuthorizationSecretGeneratorCustomResource7B87C244",
               },
-              "\\"
+              "\\",
+  \\"httpHeaders\\": []
 }",
             ],
           ],
@@ -1788,7 +1792,8 @@ Object {
                   "UserPoolClient.ClientSecret",
                 ],
               },
-              "\\"
+              "\\",
+  \\"httpHeaders\\": []
 }",
             ],
           ],
@@ -1875,7 +1880,8 @@ Object {
                   "UserPoolClient.ClientSecret",
                 ],
               },
-              "\\"
+              "\\",
+  \\"httpHeaders\\": []
 }",
             ],
           ],
@@ -1962,7 +1968,8 @@ Object {
                   "UserPoolClient.ClientSecret",
                 ],
               },
-              "\\"
+              "\\",
+  \\"httpHeaders\\": []
 }",
             ],
           ],
@@ -2049,7 +2056,8 @@ Object {
                   "UserPoolClient.ClientSecret",
                 ],
               },
-              "\\"
+              "\\",
+  \\"httpHeaders\\": []
 }",
             ],
           ],

packages/cdk-cloudfront-authorization/src/auth-flow.ts

@@ -19,6 +19,7 @@ export interface AuthFlowProps {
   readonly cookieSettings: Record<string, string>;
   readonly nonceSigningSecret: string;
   readonly clientSecret?: string;
+  readonly httpHeaders?: Record<string, string>;
 }
 
 export class AuthFlow extends Construct {
@@ -44,6 +45,7 @@ export class AuthFlow extends Construct {
       cookieSettings: props.cookieSettings,
       nonceSigningSecret: props.nonceSigningSecret,
       clientSecret: props.clientSecret,
+      httpHeaders: props.httpHeaders ?? [],
     };
 
     this.checkAuth = new EdgeFunction(this, 'CheckAuth', {

packages/cdk-cloudfront-authorization/src/authorizations.ts

@@ -44,6 +44,7 @@ export interface AuthorizationProps {
   readonly oauthScopes?: aws_cognito.OAuthScope[];
   readonly cookieSettings?: Record<string, string>;
   readonly identityProviders?: aws_cognito.UserPoolClientIdentityProvider[];
+  readonly httpHeaders?: Record<string, string> | undefined;
 }
 
 export abstract class Authorization extends Construct {
@@ -59,6 +60,7 @@ export abstract class Authorization extends Construct {
   protected readonly cognitoAuthDomain: string;
   protected readonly identityProviders: aws_cognito.UserPoolClientIdentityProvider[];
   protected readonly responseHeaderPolicy: aws_cloudfront.IResponseHeadersPolicy;
+  protected readonly httpHeaders: Record<string, string> | undefined;
 
   constructor(scope: Construct, id: string, props: AuthorizationProps) {
     super(scope, id);
@@ -126,6 +128,8 @@ export abstract class Authorization extends Construct {
 
     this.cognitoAuthDomain = this.retrieveCognitoAuthDomain();
 
+    this.httpHeaders = props.httpHeaders;
+
     this.authFlow = this.createAuthFlow(props.logLevel ?? LogLevel.WARN);
   }
 
@@ -297,6 +301,7 @@ export class StaticSiteAuthorization extends Authorization implements IStaticSit
         refreshToken: 'Path=/; Secure; HttpOnly; SameSite=Lax',
         nonce: 'Path=/; Secure; HttpOnly; SameSite=Lax',
       },
+      httpHeaders: this.httpHeaders,
     });
   }
 

packages/cdk-cloudfront-authorization/src/lambdas/user-pool-domain/index.ts

@@ -7,8 +7,9 @@ import type {
 } from 'aws-lambda';
 import axios from 'axios';
 
+const awsRegion = process.env.AWS_REGION || process.env.AWS_DEFAULT_REGION;
 const COGNITO_CLIENT = new CognitoIdentityProvider({
-  region: process.env.AWS_REGION || process.env.AWS_DEFAULT_REGION,
+  region: awsRegion,
 });
 
 async function ensureCognitoUserPoolDomain(userPoolId: string): Promise<string> {
@@ -28,7 +29,7 @@ async function ensureCognitoUserPoolDomain(userPoolId: string): Promise<string>
     throw new Error('Cognito auth domain is missing! Either a domain prefix or a custom domain must be configured.');
   }
 
-  return userPool.CustomDomain ?? `${userPool.Domain}.auth.${COGNITO_CLIENT.config.region}.amazoncognito.com`;
+  return userPool.CustomDomain ?? `${userPool.Domain}.auth.${awsRegion}.amazoncognito.com`;
 }
 
 export const handler: CloudFormationCustomResourceHandler = async (event) => {

packages/cdk-lambda-at-edge-pattern/src/lambdas/with-configuration/index.ts

@@ -1,6 +1,6 @@
 import { mkdtempSync, writeFileSync } from 'fs';
 import { resolve } from 'path';
-import { GetFunctionCommand, LambdaClient, UpdateFunctionCodeCommand } from '@aws-sdk/client-lambda';
+import { GetFunctionCommand, LambdaClient, UpdateFunctionCodeCommand, waitUntilFunctionActiveV2 } from '@aws-sdk/client-lambda';
 import Zip from 'adm-zip';
 import type {
   CloudFormationCustomResourceCreateEvent,
@@ -26,6 +26,7 @@ interface WithConfiguration {
 const updateLambdaCode = async (
   event: CloudFormationCustomResourceCreateEvent | CloudFormationCustomResourceUpdateEvent,
 ): Promise<ResourceHandlerReturn> => {
+  console.log(`Recieved event: ${event}`);
   const { region, functionName, configuration } = camelizeKeys<
     WithConfiguration,
     CloudFormationCustomResourceEventCommon['ResourceProperties']
@@ -72,15 +73,48 @@ const updateLambdaCode = async (
       Publish: true,
     }),
   );
-
-  return {
+
+  let responseDetails = {
+    responseStatus: 'FAILED',
+    reason: 'Internal Error',
     physicalResourceId: functionName,
-    responseData: {
-      CodeSha256: codeSha256,
-      Version: version,
-      FunctionArn: functionArn,
-    },
+    responseData: {},
   };
+
+  // wait for functions to go active before proceeding because if an edge function isn't active updates to the
+  // related distribution fail
+  console.log(`Waiting for ${functionName} to go active...`);
+  waitUntilFunctionActiveV2(
+    {client: lambda, maxWaitTime: 60},
+    {FunctionName: functionName}
+  )
+  .then((result=>{
+    console.log(`Lambda function ${functionName} is active. Result: ${result}`);
+    responseDetails = {
+      physicalResourceId: functionName,
+      responseStatus: 'SUCCESS',
+      reason: `Lambda function ${functionName} is active.`,
+      responseData: {
+        CodeSha256: codeSha256,
+        Version: version,
+        FunctionArn: functionArn,
+      },
+    };
+    setTimeout(()=>{
+      console.log(`Waiting 5s for lambda status propagation...`);
+    }, 5000);
+  }))
+  .catch((err=>{
+    console.log(`Error waiting for ${functionName} to go active: ${err}`);
+    responseDetails = {
+      responseStatus: 'FAILED',
+      physicalResourceId: functionName,
+      reason: `Encountered error while waiting for ${functionName} to go active: ${err}`,
+      responseData: {}
+    };
+  }));
+
+  return responseDetails;
 };
 
 const handleCreate: OnCreateHandler = async (event): Promise<ResourceHandlerReturn> => updateLambdaCode(event);